Merge pull request #5524 from hantek/allocdiag

theano/gpuarray/subtensor.py

@@ -21,7 +21,6 @@ from .type import GpuArrayType, gpu_context_type
 from .basic_ops import (as_gpuarray_variable, HideC, GpuKernelBase, Kernel,
                         infer_context_name, gpu_contiguous)
 
-
 iadd_reg = {}
 
 
@@ -1080,7 +1079,7 @@ __device__ ga_half atomicExch(ga_half *addr, ga_half val) {
         """ % locals()
 
 
-class GpuDiagonal(Subtensor):
+class GpuExtractDiag(Op):
     __props__ = ("offset", "axis1", "axis2", "view")
 
     def __init__(self, offset=0, axis1=0, axis2=1, view=False):
@@ -1102,9 +1101,7 @@ class GpuDiagonal(Subtensor):
         broadcastable = x.broadcastable[:axis_small] + \
             x.broadcastable[axis_small + 1:axis_large] + \
             x.broadcastable[axis_large + 1:] + (False,)
-        return gof.Apply(self, [x], [x.type.__class__(
-            dtype=x.dtype,
-            broadcastable=broadcastable)()])
+        return gof.Apply(self, [x], [x.type.clone(broadcastable=broadcastable)()])
 
     def perform(self, node, inputs, outputs):
         (x,) = inputs
@@ -1183,3 +1180,42 @@ class GpuDiagonal(Subtensor):
             diag_size = T.minimum(dim1, dim2)
         out_shape.append(diag_size)
         return [tuple(out_shape)]
+
+
+class GpuAllocDiag(Op):
+    __props__ = ("offset",)
+
+    def __init__(self, offset=0):
+        self.offset = offset
+
+    def make_node(self, _x):
+        ctx_name = infer_context_name(_x)
+        x = as_gpuarray_variable(_x, ctx_name)
+
+        if x.ndim != 1:
+            raise ValueError('AllocDiag argument must be a vector!', x)
+
+        return gof.Apply(self, [x], [x.type.clone(broadcastable=(False, False))()])
+
+    def perform(self, node, inputs, outputs):
+        (x,) = inputs
+        (z,) = outputs
+
+        dim = x.shape[0] + abs(self.offset)
+        z[0] = gpuarray.zeros((dim, dim), dtype=x.dtype, context=x.context)
+
+        if self.offset <= 0:  # diag in the lower triangle
+            diag_z = z[0][-self.offset, :(dim + self.offset)]
+        else:  # diag in the upper triangle
+            diag_z = z[0][:(dim - self.offset), self.offset]
+        diag_z.strides = (sum(z[0].strides),)
+
+        diag_z[:] = x[:]
+
+    def grad(self, inputs, gout):
+        (gz,) = gout
+        return [GpuExtractDiag(offset=self.offset, axis1=0, axis2=1)(gz)]
+
+    def infer_shape(self, node, shapes):
+        dim = shapes[0][0] + abs(self.offset)
+        return [[dim, dim]]

theano/gpuarray/tests/test_subtensor.py

@@ -15,7 +15,8 @@ from ..subtensor import (GpuIncSubtensor, GpuSubtensor,
                          GpuAdvancedSubtensor,
                          GpuAdvancedIncSubtensor1,
                          GpuAdvancedIncSubtensor1_dev20,
-                         GpuDiagonal)
+                         GpuExtractDiag,
+                         GpuAllocDiag)
 from ..type import gpuarray_shared_constructor
 
 from .config import mode_with_gpu
@@ -191,15 +192,15 @@ def test_adv_subtensor():
     assert np.allclose(rval, rep)
 
 
-class test_gpudiagonal(unittest.TestCase):
+class test_gpuextractdiag(unittest.TestCase):
     def test_matrix(self):
         x = tensor.matrix()
         np_x = np.arange(77).reshape(7, 11).astype(theano.config.floatX)
-        fn = theano.function([x], GpuDiagonal()(x), mode=mode_with_gpu)
+        fn = theano.function([x], GpuExtractDiag()(x), mode=mode_with_gpu)
         assert np.allclose(fn(np_x), np_x.diagonal())
-        fn = theano.function([x], GpuDiagonal(2)(x), mode=mode_with_gpu)
+        fn = theano.function([x], GpuExtractDiag(2)(x), mode=mode_with_gpu)
         assert np.allclose(fn(np_x), np_x.diagonal(2))
-        fn = theano.function([x], GpuDiagonal(-3)(x), mode=mode_with_gpu)
+        fn = theano.function([x], GpuExtractDiag(-3)(x), mode=mode_with_gpu)
         assert np.allclose(fn(np_x), np_x.diagonal(-3))
 
     def test_tensor(self):
@@ -210,5 +211,65 @@ class test_gpudiagonal(unittest.TestCase):
                 (-3, 1, 0), (-2, 2, 0), (3, 3, 0), (-1, 3, 2),
                 (2, 2, 3), (-1, 2, 1), (1, 3, 1), (-1, 1, 3)]:
             assert np.allclose(
-                GpuDiagonal(offset, axis1, axis2)(x).eval({x: np_x}),
+                GpuExtractDiag(offset, axis1, axis2)(x).eval({x: np_x}),
                 np_x.diagonal(offset, axis1, axis2))
+
+
+class test_gpuallocdiag(unittest.TestCase):
+    def test_matrix(self):
+        x = tensor.vector()
+        np_x = np.arange(7).astype(theano.config.floatX)
+        fn = theano.function([x], GpuAllocDiag()(x), mode=mode_with_gpu)
+        assert np.allclose(fn(np_x), np.diag(np_x))
+        fn = theano.function([x], GpuAllocDiag(2)(x), mode=mode_with_gpu)
+        assert np.allclose(fn(np_x), np.diag(np_x, 2))
+        fn = theano.function([x], GpuAllocDiag(-3)(x), mode=mode_with_gpu)
+        assert np.allclose(fn(np_x), np.diag(np_x, -3))
+
+    def test_grad(self):
+        x = tensor.vector()
+        np_x = np.random.randn(7).astype(theano.config.floatX)
+
+        # offset = 0 case:
+        mtx_x = GpuAllocDiag()(x)
+        sum_mtx_x = tensor.sum(mtx_x)
+        grad_x = tensor.grad(sum_mtx_x, x)
+        grad_mtx_x = tensor.grad(sum_mtx_x, mtx_x)
+
+        fn_grad_x = theano.function([x], grad_x)
+        fn_grad_mtx_x = theano.function([x], grad_mtx_x)
+
+        computed_grad_x = fn_grad_x(np_x)
+        computed_grad_mtx_x = fn_grad_mtx_x(np_x)
+        true_grad_x = np.diagonal(computed_grad_mtx_x, 0)
+        assert np.allclose(computed_grad_x, true_grad_x)
+
+        # offset > 0 case:
+        mtx_x = GpuAllocDiag(2)(x)
+        sum_mtx_x = tensor.sum(mtx_x)
+        grad_x = tensor.grad(sum_mtx_x, x)
+        grad_mtx_x = tensor.grad(sum_mtx_x, mtx_x)
+
+        fn_grad_x = theano.function([x], grad_x)
+        fn_grad_mtx_x = theano.function([x], grad_mtx_x)
+
+        computed_grad_x = fn_grad_x(np_x)
+        computed_grad_mtx_x = fn_grad_mtx_x(np_x)
+        true_grad_x = np.diagonal(computed_grad_mtx_x, 2)
+        assert np.allclose(computed_grad_x, true_grad_x)
+
+        # offset < 0 case:
+        mtx_x = GpuAllocDiag(-3)(x)
+        sum_mtx_x = tensor.sum(mtx_x)
+        grad_x = tensor.grad(sum_mtx_x, x)
+        grad_mtx_x = tensor.grad(sum_mtx_x, mtx_x)
+
+        fn_grad_x = theano.function([x], grad_x)
+        fn_grad_mtx_x = theano.function([x], grad_mtx_x)
+
+        computed_grad_x = fn_grad_x(np_x)
+        computed_grad_mtx_x = fn_grad_mtx_x(np_x)
+        true_grad_x = np.diagonal(computed_grad_mtx_x, -3)
+        assert np.allclose(computed_grad_x, true_grad_x)
+
+        # assert

theano/tensor/basic.py

@@ -6285,6 +6285,17 @@ class ExtractDiag(Op):
     def grad(self, inputs, gout):
         (x,) = inputs
         (gz,) = gout
+
+        if x.ndim == 2:
+            # The following code is moved from tensor.nlinalg.ExtractDiag, only
+            # works for matrices.
+            x = theano.tensor.zeros_like(x)
+            xdiag = theano.tensor.AllocDiag(offset=self.offset)(gz)
+            return [theano.tensor.set_subtensor(
+                x[:xdiag.shape[0], :xdiag.shape[1]], xdiag)]
+        else:
+            warnings.warn("gradient of theano.tensor.nlinalg.ExtractDiag only"
+                          "works for matrices.")
             return [grad_not_implemented(self, 0, x)]
 
     def infer_shape(self, node, shapes):
@@ -6306,42 +6317,108 @@ class ExtractDiag(Op):
 
 
 def diagonal(a, offset=0, axis1=0, axis2=1):
-    if (offset, axis1, axis2) == (0, 0, 1):
-        return theano.tensor.nlinalg.extract_diag(a)
+    """
+    A helper function for `theano.tensor.ExtractDiag`. It accepts tensor with
+    `ndim >= 2` as input. The name `diagonal` is just meant to keep it
+    consistent with numpy.
+
+    Parameters
+    ----------
+    a : symbolic tensor
+    offset : int
+        offset
+    axis1 : int
+    axis2 : int
+
+    Returns
+    -------
+    tensor : symbolic tensor
+
+    """
     return ExtractDiag(offset, axis1, axis2)(a)
 
 
-class Diag(Op):
+class AllocDiag(Op):
+    """
+    An op that copies a vector to the diagonal of an empty matrix. It does the
+    inverse of ExtractDiag.
 
-    __props__ = ()
+    Usage: T.AllocDiag()(x)
+
+    `x` should be a tensor vector. The parenthesis in the front should indicate
+    which main diagonal the vector value goes into. By default it is set to
+    `0`, which corresponds to setting the values of x to the main diagonal in
+    the returned matrix.
+
+    Parameters
+    ----------
+    offset : int
+        Indicates which diagonal to put `x` into. Defaults to `0`.
+
+    x: symbolic vector
+        A tensor vector consists of diagonal values.
+
+    Returns
+    -------
+    tensor : symbolic tenstor
+        A tensor with passed vector values at its corresponding diagonal.
+
+    """
+
+    __props__ = ("offset", )
+    default_offset = 0
+
+    def __init__(self, offset=0):
+        if numpy_diagonal_return_view:
+            self.view_map = {0: [0]}
+        self.offset = offset
 
     def make_node(self, diag):
         diag = as_tensor_variable(diag)
         if diag.type.ndim != 1:
             raise TypeError('data argument must be a vector', diag.type)
-
         return Apply(self, [diag], [matrix(dtype=diag.dtype)])
 
     def perform(self, node, inputs, outputs):
         (z,) = outputs
-        z[0] = numpy.diag(inputs[0])
+        z[0] = numpy.diag(inputs[0], self.offset)
 
     def grad(self, inputs, gout):
         (gz,) = gout
-        return [diagonal(gz)]
+        return [diagonal(gz, offset=self.offset, axis1=0, axis2=1)]
 
     def infer_shape(self, nodes, shapes):
         return [(shapes[0][0],) * 2]
 
 
 def diag(v, k=0):
+    """
+    A helper function for two ops: `theano.tensor.ExtractDiag` and
+    `theano.tensor.AllocDiag`. The name `diag` is meant to keep it consistent
+    with numpy. It both accepts tensor vector and tensor matrix.
+    While the passed tensor variable `v` has `v.ndim>=2`, it builds a
+    `ExtractDiag` instance, and returns a vector with its entries equal to
+    `v`'s main diagonal; otherwise if `v.ndim` is `1`, it builds an `AllocDiag`
+    instance, and returns a matrix with `v` at its k-th diaogonal.
+
+    Parameters
+    ----------
+    v : symbolic tensor
+    k : int
+        offset
+
+    Returns
+    -------
+    tensor : symbolic tensor
+
+    """
+
     if v.ndim == 1:
-        assert k == 0, "diagonals other than main are not implemented"
-        return Diag()(v)
-    elif v.ndim == 2:
-        return diagonal(v, k)
+        return AllocDiag(k)(v)
+    elif v.ndim >= 2:
+        return diagonal(v, offset=k)
     else:
-        raise ValueError("Input must be 1- or 2-d.")
+        raise ValueError("Input must has v.ndim >= 1.")
 
 
 def stacklists(arg):

theano/tensor/nlinalg.py

 from __future__ import absolute_import, print_function, division
 
 import logging
-
+import warnings
 import numpy
 from six.moves import xrange
 
@@ -145,6 +145,10 @@ class AllocDiag(Op):
     __props__ = ()
 
     def make_node(self, _x):
+        warnings.warn("DeprecationWarning: theano.tensor.nlinalg.AllocDiag"
+                      "is deprecated, please use theano.tensor.AllocDiag"
+                      "instead.",
+                      category=DeprecationWarning)
         x = as_tensor_variable(_x)
         if x.type.ndim != 1:
             raise TypeError('AllocDiag only works on vectors', _x)
@@ -184,6 +188,10 @@ class ExtractDiag(Op):
             self.view_map = {0: [0]}
 
     def make_node(self, _x):
+        warnings.warn("DeprecationWarning: theano.tensor.nlinalg.ExtractDiag"
+                      "is deprecated, please use theano.tensor.ExtractDiag"
+                      "instead.",
+                      category=DeprecationWarning)
         if not isinstance(_x, theano.Variable):
             x = as_tensor_variable(_x)
         else: