import math
import numpy
from pycuda import gpuarray

def _sqnorm(x):
    if type(x) == gpuarray.GPUArray:
        return float(gpuarray.dot(x, x).get())
    x = x.ravel()
    return float(x.dot(x))

class Optimizer(object):
    """Optimizers' base class."""

    def setup(self, params_grads):
        self.tuples = [(p, g, self.init_state(p, g))
                       for p, g in zip(*params_grads)]
        self.t = 0

    def init_state(self, param, grad):
        if isinstance(param, gpuarray.GPUArray):
            return self.init_state_gpu(param, grad)
        return self.init_state_cpu(param, grad)

    def init_state_cpu(self, param, grad):
        """Initialize state on CPU. Child class using state should override it."""
        return None

    def init_state_gpu(self, param, grad):
        """Initialize state on GPU. Child class using state should override it."""
        return None

    def zero_grads(self):
        """Set gradients zero."""
        for _, g, _ in self.tuples:
            g.fill(0)

    def compute_grads_norm(self):
        """Compute norm of the gradient."""
        sqnorm = 0
        for _, g, _ in self.tuples:
            sqnorm += _sqnorm(g)
        return math.sqrt(sqnorm)

    def clip_grads(self, maxnorm):
        """Clip norm of the gradient."""
        norm = self.compute_grads_norm()
        if norm > maxnorm:
            ratio = maxnorm / norm
            for _, g, _ in self.tuples:
                g *= ratio

    def weight_decay(self, decay):
        """Apply weight decay."""
        for p, g, _ in self.tuples:
            g -= decay * p

    def update(self):
        self.t += 1
        for p, g, s in self.tuples:
            self.update_one(p, g, s)

    def update_one(self, param, grad, state):
        if isinstance(param, gpuarray.GPUArray):
            self.update_one_gpu(param, grad, state)
        else:
            self.update_one_cpu(param, grad, state)

    def update_one_cpu(self, param, grad, state):
        raise NotImplementedError()

    def update_one_gpu(self, param, grad, state):
        raise NotImplementedError()

@@ -45,16 +45,28 @@
 gpuarray
+%0Aimport cuda
 %0A%0Adef _s
@@ -86,27 +86,26 @@
  if 
-type(x) == gpuarray
+isinstance(x, cuda
 .GPU
@@ -109,16 +109,17 @@
 GPUArray
+)
 :%0A      
@@ -999,32 +999,172 @@
 in self.tuples:%0A
+            if isinstance(g, cuda.GPUArray):%0A                with cuda.using_device(g):%0A                    g.fill(0)%0A            else:%0A    
             g.fi
@@ -1240,32 +1240,112 @@
 he gradient.%22%22%22%0A
+        # TODO(beam2d): Make it asynchronous to CPU when gradients exist on GPU%0A
         sqnorm =
@@ -1794,32 +1794,347 @@
 in self.tuples:%0A
+            if isinstance(p, cuda.GPUArray):%0A                with cuda.using_device(p):%0A                    cuda.elementwise('float* g, const float* p, float decay',%0A                                     'g%5Bi%5D -= decay * p%5Bi%5D',%0A                                     'weight_decay')(g, p, decay)%0A            else:%0A    
             g -=
@@ -2332,32 +2332,28 @@
 ance(param, 
-gpuarray
+cuda
 .GPUArray):%0A
@@ -2344,32 +2344,79 @@
 cuda.GPUArray):%0A
+            with cuda.using_device(param):%0A    
             self

Support model parallel