Merge pull request #6386 from shawntan/issue-5633

theano/tensor/basic.py

@@ -6333,18 +6333,57 @@ del x
 
 
 class ExtractDiag(Op):
-    """Return specified diagonals.
+    """
+    Return specified diagonals.
+
+    If x is 2-D, returns the diagonal of x with the given offset,
+    i.e., the collection of elements of the form x[i, i+offset].
+    If x has more than two dimensions, then the axes specified by
+    axis1 and axis2 are used to determine the 2-D sub-array whose
+    diagonal is returned. The shape of the resulting array can be
+    determined by removing axis1 and axis2 and appending an index
+    to the right equal to the size of the resulting diagonals.
 
     Parameters
     ----------
-    x
-        A tensor variable with x.ndim >= 2.
+    x: A tensor variable with x.ndim >= 2.
+
+    offset: Offset of the diagonal from the main diagonal.
+        Can be positive or negative.
+        Defaults to main diagonal (0).
+
+    axis1: Axis to be used as the first axis of the 2-D
+        sub-arrays from which the diagonals should be taken.
+        Defaults to first axis (0).
+
+    axis2: Axis to be used as the second axis of the 2-D
+        sub-arrays from which the diagonals should be taken.
+        Defaults to second axis (1).
+
+
 
     Returns
     -------
-    vector
-        A vector representing the diagonal elements.
+    array_of_diagonals:
+        If x is 2-D, a 1-D array of the same type as a
+        containing the diagonal is returned.
+        If the dimension of x is greater than two, then an
+        array of diagonals is returned, "packed" from left-most
+        dimension to right-most (e.g., if x is 3-D, then the
+        diagonals are "packed" along rows).
+
+
+
+    Raises
+    ------
+    ValueError
+        If the dimension of x is less than 2.
+
 
+    See Also
+    --------
+    numpy.diagonal:
+        https://docs.scipy.org/doc/numpy-dev/reference/generated/numpy.diagonal.html
     """
     __props__ = ("offset", "axis1", "axis2", "view")
 
@@ -6385,14 +6424,12 @@ class ExtractDiag(Op):
         (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"
+            warnings.warn("gradient of theano.tensor.basic.ExtractDiag only"
                           "works for matrices.")
             return [grad_not_implemented(self, 0, x)]
 
@@ -6413,6 +6450,26 @@ class ExtractDiag(Op):
         out_shape.append(diag_size)
         return [tuple(out_shape)]
 
+    def __setstate__(self, state):
+        self.__dict__.update(state)
+        if self.view and not numpy_diagonal_return_view:
+            warnings.warn("View will forced to False. ExtractDiag property view is "
+                          "set to True but numpy version %s and prior versions of "
+                          "numpy.diagonal() do not return a view. Update "
+                          "numpy to use ExtractDiag(view=True)" %
+                          np.version.version)
+            self.view = False
+
+        if self.view:
+            self.view_map = {0: [0]}
+
+        if "offset" not in state:
+            self.offset = 0
+        if "axis1" not in state:
+            self.axis1 = 0
+        if "axis2" not in state:
+            self.axis2 = 1
+
 
 def diagonal(a, offset=0, axis1=0, axis2=1):
     """

theano/tensor/nlinalg.py

@@ -10,7 +10,7 @@ from theano.tensor import as_tensor_variable
 from theano.gof import Op, Apply
 from theano.gradient import DisconnectedType
 from theano.tensor import basic as tensor
-
+from theano.tensor.basic import ExtractDiag
 logger = logging.getLogger(__name__)
 
 
@@ -193,75 +193,6 @@ class AllocDiag(Op):
         return [(x_s[0], x_s[0])]
 
 alloc_diag = AllocDiag()
-
-
-class ExtractDiag(Op):
-    """Return the diagonal of a matrix.
-
-    Notes
-    -----
-    Works on the GPU.
-
-    """
-
-    __props__ = ("view",)
-
-    def __init__(self, view=False):
-        self.view = view
-        if self.view:
-            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:
-            x = _x
-
-        if x.type.ndim != 2:
-            raise TypeError('ExtractDiag only works on matrices', _x)
-        y = x.type.clone(broadcastable=(False,))()
-        return Apply(self, [x], [y])
-
-    def perform(self, node, ins, outs):
-        """ For some reason numpy.diag(x) is really slow, so we
-        implemented our own. """
-        x, = ins
-        z, = outs
-        # zero-dimensional matrices ...
-        if x.shape[0] == 0 or x.shape[1] == 0:
-            z[0] = node.outputs[0].type.value_zeros((0,))
-            return
-
-        if x.shape[0] < x.shape[1]:
-            rval = x[:, 0]
-        else:
-            rval = x[0]
-
-        rval.strides = (x.strides[0] + x.strides[1],)
-        if self.view:
-            z[0] = rval
-        else:
-            z[0] = rval.copy()
-
-    def __str__(self):
-        return 'ExtractDiag{view=%s}' % self.view
-
-    def grad(self, inputs, g_outputs):
-        x = theano.tensor.zeros_like(inputs[0])
-        xdiag = alloc_diag(g_outputs[0])
-        return [theano.tensor.set_subtensor(
-            x[:xdiag.shape[0], :xdiag.shape[1]],
-            xdiag)]
-
-    def infer_shape(self, node, shapes):
-        x_s, = shapes
-        shp = theano.tensor.min(node.inputs[0].shape)
-        return [(shp,)]
-
 extract_diag = ExtractDiag()
 # TODO: optimization to insert ExtractDiag with view=True
 

theano/tensor/tests/test_nlinalg.py

@@ -349,15 +349,6 @@ class test_diag(unittest.TestCase):
         y = extract_diag(x)
         assert y.owner.op.__class__ == ExtractDiag
 
-        # other types should raise error
-        x = theano.tensor.tensor3()
-        ok = False
-        try:
-            y = extract_diag(x)
-        except TypeError:
-            ok = True
-        assert ok
-
     # not testing the view=True case since it is not used anywhere.
     def test_extract_diag(self):
         rng = np.random.RandomState(utt.fetch_seed())
@@ -384,6 +375,8 @@ class test_diag(unittest.TestCase):
             extract_diag(xx)
         except TypeError:
             ok = True
+        except ValueError:
+            ok = True
         assert ok
 
         # Test infer_shape
@@ -429,6 +422,9 @@ def test_trace():
         trace(xx)
     except TypeError:
         ok = True
+    except ValueError:
+        ok = True
+
     assert ok