extract diag

theano/gpuarray/subtensor.py

@@ -7,6 +7,8 @@ from six import integer_types
 from six.moves import StringIO
 
 from theano import tensor, gof, Op
+from theano.gradient import grad_not_implemented
+import theano.tensor.clip, theano.tensor.minimum
 from theano.tensor.subtensor import IncSubtensor, Subtensor, get_idx_list
 
 try:
@@ -1078,7 +1080,7 @@ __device__ ga_half atomicExch(ga_half *addr, ga_half val) {
         """ % locals()
 
 
-class GpuDiagonal(GpuOp):
+class GpuDiagonal(Subtensor):
     __props__ = ("offset", "axis1", "axis2", "view")
 
     def __init__(self, offset=0, axis1=0, axis2=1, view=False):
@@ -1090,7 +1092,8 @@ class GpuDiagonal(GpuOp):
         self.axis2 = axis2
 
     def make_node(self, _x):
-        x = as_cuda_ndarray_variable(_x)
+        ctx_name = infer_context_name(_x)
+        x = as_gpuarray_variable(_x, ctx_name)
 
         if x.ndim < 2:
             raise ValueError('Diagonal needs an input with 2 or more '
@@ -1099,7 +1102,7 @@ class GpuDiagonal(GpuOp):
         broadcastable = x.broadcastable[:axis_small] + \
             x.broadcastable[axis_small + 1:axis_large] + \
             x.broadcastable[axis_large + 1:] + (False,)
-        return Apply(self, [x], [x.type.__class__(
+        return gof.Apply(self, [x], [x.type.__class__(
             dtype=x.dtype,
             broadcastable=broadcastable)()])
 
@@ -1146,7 +1149,8 @@ class GpuDiagonal(GpuOp):
         new_dim_order = tuple(new_dim_order[:stride_axis] +
                               new_dim_order[stride_axis + 1:] +
                               [stride_axis, ])
-        rval = cuda_ndarray.cuda_ndarray.dimshuffle(x[slicer], new_dim_order)
+        # rval = cuda_ndarray.cuda_ndarray.dimshuffle(x[slicer], new_dim_order)
+        rval = x[slicer].transpose(new_dim_order)
 
         # step 3) modify the strides in the last axis, such that rval becomes
         # a view on the diagonal.

theano/gpuarray/tests/test_subtensor.py

 from __future__ import absolute_import, print_function, division
 import numpy
+import unittest
 
 import theano
 from theano import tensor
@@ -11,7 +12,8 @@ from ..elemwise import GpuDimShuffle
 from ..subtensor import (GpuIncSubtensor, GpuSubtensor,
                          GpuAdvancedSubtensor1,
                          GpuAdvancedSubtensor,
-                         GpuAdvancedIncSubtensor1)
+                         GpuAdvancedIncSubtensor1,
+                         GpuDiagonal)
 from ..type import gpuarray_shared_constructor
 
 from .config import mode_with_gpu
@@ -132,11 +134,11 @@ class test_gpudiagonal(unittest.TestCase):
     def test_matrix(self):
         x = cuda.fmatrix()
         np_x = numpy.arange(77).reshape(7, 11).astype('float32')
-        fn = theano.function([x], B.GpuDiagonal()(x), mode=mode_with_gpu)
+        fn = theano.function([x], GpuDiagonal()(x), mode=mode_with_gpu)
         assert numpy.allclose(fn(np_x), np_x.diagonal())
-        fn = theano.function([x], B.GpuDiagonal(2)(x), mode=mode_with_gpu)
+        fn = theano.function([x], GpuDiagonal(2)(x), mode=mode_with_gpu)
         assert numpy.allclose(fn(np_x), np_x.diagonal(2))
-        fn = theano.function([x], B.GpuDiagonal(-3)(x), mode=mode_with_gpu)
+        fn = theano.function([x], GpuDiagonal(-3)(x), mode=mode_with_gpu)
         assert numpy.allclose(fn(np_x), np_x.diagonal(-3))
 
     def test_tensor(self):
@@ -147,6 +149,5 @@ 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 numpy.allclose(
-                B.GpuDiagonal(offset, axis1, axis2)(x).eval({x: np_x}),
+                GpuDiagonal(offset, axis1, axis2)(x).eval({x: np_x}),
                 np_x.diagonal(offset, axis1, axis2))
-

theano/tensor/basic.py

@@ -6098,7 +6098,7 @@ 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)
-    return Diagonal(offset, axis1, axis2)(a)
+    return ExtractDiag(offset, axis1, axis2)(a)
 
 
 class Diag(Op):