L2 Normalize layer for Chainer

l2norm.py

class L2Normalize(Function):
    
    def forward(self, inputs):
        xp = cuda.get_array_module(*inputs)
        x, = inputs
        if type(x) is cuda.ndarray:
            Xnp = cuda.cupy.asnumpy(x)
            self._norm = xp.asarray(xp.expand_dims(np.linalg.norm(Xnp, ord=2, axis=1), axis=1), dtype=xp.float32)
        else:
            self._norm = xp.expand_dims(np.linalg.norm(x, ord=2, axis=1), axis=1)
        z = xp.divide(x, self._norm)
        return z,

    def backward(self, inputs, grad_outputs):
        #xp = cuda.get_array_module(*inputs)
        x, = inputs
        gz, = grad_outputs
        gx = None # how to compute gradient here?
        return gx,
        
def l2_normalize(X):
    return L2Normalize()(X)