import numpy
from chainer import cuda, Optimizer

class AdaDelta(Optimizer):
    """Zeiler's ADADELTA.

    See: http://www.matthewzeiler.com/pubs/googleTR2012/googleTR2012.pdf

    """
    def __init__(self, rho=0.95, eps=1e-6):
        self.rho = rho
        self.eps = eps

    def init_state_cpu(self, param, grad):
        return numpy.zeros_like(param), numpy.zeros_like(param)

    def init_state_gpu(self, param, grad):
        return cuda.zeros_like(param), cuda.zeros_like(param)

    def update_one_cpu(self, param, grad, state):
        msg, msdx = state
        msg *= self.rho
        msg += (1 - self.rho) * grad * grad
        dx = - numpy.sqrt((msdx + self.eps) / (msg + self.eps)) * grad
        msdx *= self.rho
        msdx += (1 - self.rho) * dx * dx
        param += dx

    def update_one_gpu(self, param, grad, ms):
        cuda.elementwise(
            '''float* param, const float* grad, float* msg, float* msdx,
               float rho, float eps''',
            '''msg[i]   = rho * msg[i]  + (1 - rho) * grad[i] * grad[i];
               float dx = - sqrtf((msdx + eps) / (msg + eps)) * grad[i];
               msdx[i]  = rho * msdx[i] + (1 - rho) * dx * dx;
               param[i] += dx;''',
            'adadelta')(param, grad, msg, msdx, self.rho, self.eps)

@@ -640,18 +640,16 @@
     dx =
- -
  numpy.s
@@ -689,32 +689,32 @@
 lf.eps)) * grad%0A
+
         msdx *= 
@@ -777,17 +777,17 @@
   param 
-+
+-
 = dx%0A%0A  
@@ -951,16 +951,26 @@
   float 
+one_minus_
 rho, flo
@@ -1008,38 +1008,28 @@
 %5D   
++
 = 
-rho * msg%5Bi%5D  + (1 - 
+one_minus_
 rho
-)
  * 
+(
 grad
@@ -1033,32 +1033,42 @@
 rad%5Bi%5D * grad%5Bi%5D
+ - msg%5Bi%5D)
 ;%0A              
@@ -1081,11 +1081,10 @@
  dx 
+ 
 =
- -
  sqr
@@ -1091,16 +1091,19 @@
 tf((msdx
+%5Bi%5D
  + eps) 
@@ -1108,16 +1108,19 @@
 ) / (msg
+%5Bi%5D
  + eps))
@@ -1159,45 +1159,46 @@
 i%5D  
++
 = 
-rho * msdx%5Bi%5D + (1 - rho) * dx * dx
+one_minus_rho * (dx * dx - msdx%5Bi%5D)
 ;%0A  
@@ -1219,17 +1219,17 @@
 aram%5Bi%5D 
-+
+-
 = dx;'''
@@ -1277,16 +1277,20 @@
 g, msdx,
+ 1 -
  self.rh

Fix and slightly improve AdaDelta