modified: ../__init__.py

theano/tensor/__init__.py

@@ -62,4 +62,5 @@ from theano.gradient import Rop, Lop, grad, numeric_grad, verify_grad, \
 
 from theano.tensor.sort import sort, argsort
 from theano.tensor.extra_ops import (DiffOp, bincount, squeeze,
-                       repeat, bartlett, fill_diagonal, cumsum, cumprod)
+                       repeat, bartlett, fill_diagonal, fill_diagonal_offset,
+                       cumsum, cumprod)

theano/tensor/extra_ops.py

@@ -725,3 +725,100 @@ def fill_diagonal(a, val):
     .. versionadded:: 0.6
     """
     return fill_diagonal_(a, val)
+
+
+
+# Offset version of fill_diagonal
+class FillDiagonalOffset(gof.Op):
+    # See function fill_diagonal for docstring
+    def __eq__(self, other):
+        return type(self) == type(other)
+
+    def __hash__(self):
+        return hash(type(self))
+
+    def __str__(self):
+        return self.__class__.__name__
+
+    def infer_shape(self, node, in_shapes):
+        return [in_shapes[0]]
+
+    def make_node(self, a, val, offset):
+        a = tensor.as_tensor_variable(a)
+        val = tensor.as_tensor_variable(val)
+        offset = tensor.as_tensor_variable(offset)
+        if a.ndim != 2:
+            raise TypeError('%s: first parameter must have exactly'
+                            ' two dimensions' % self.__class__.__name__)
+        elif val.ndim != 0:
+            raise TypeError('%s: second parameter must be a scalar'\
+                            % self.__class__.__name__)
+        val = tensor.cast(val, dtype=scalar.upcast(a.dtype, val.dtype))
+        if val.dtype != a.dtype:
+            raise TypeError('%s: type of second parameter must be compatible'
+                            ' with first\'s' % self.__class__.__name__)
+
+        return gof.Apply(self, [a, val, offset], [a.type()])
+
+    def perform(self, node, inputs, output_storage):
+        a = inputs[0].copy()
+        val = inputs[1]
+        offset = inputs[2]
+
+        # offset should be an integer
+        if offset % 1 != 0:
+            raise TypeError('%s: third parameter must be an integer'\
+                            % self.__class__.__name__)
+
+        # numpy.fill_diagonal up to date(including 1.6.2) have a
+        # bug for tall matrix.
+        # the offset function is only implemented for matrices
+        if offset >= 0:
+            start = offset
+        else:
+            start = - offset * a.shape[0]
+        step = a.shape[1] + 1
+        end = a.shape[1] * a.shape[1]
+        # Write the value out into the diagonal.
+        a.flat[start:end:step] = val
+
+
+        output_storage[0][0] = a
+
+    def grad(self, inp, cost_grad):
+        """
+        Note: The gradient is currently implemented for matrices
+        only.
+        """
+        a, val, offset = inp
+        grad = cost_grad[0]
+        if (a.dtype.startswith('complex')):
+            return [None, None]
+        elif a.ndim > 2:
+            raise NotImplementedError('%s: gradient is currently implemented'
+                            ' for matrices only' % self.__class__.__name__)
+        wr_a = fill_diagonal_offset(grad, 0, offset)  # valid for any number of dimensions
+        # diag is only valid for matrices
+        import theano.sandbox.linalg
+        wr_val = theano.sandbox.linalg.ops.diag(grad).sum()
+        return [wr_a, wr_val]
+fill_diagonal_offset_ = FillDiagonalOffset()
+
+
+#I create a function only to have the doc show well.
+def fill_diagonal_offset(a, val, offset):
+    """ Returns a copy of an array with all
+    elements of the main diagonal set to a specified scalar value.
+
+    :param a: Rectangular array of two dimensions.
+    :param val: Scalar value to fill the diagonal whose type must be
+        compatible with that of array 'a' (i.e. 'val' cannot be viewed
+        as an upcast of 'a').
+    :params offset : Scalar value Offset of the diagonal from the main 
+        diagonal. Can be positive or negative.
+    :return: An array identical to 'a' except that its offset diagonal
+        is filled with scalar 'val'.
+
+    Only support rectangular matrix
+    """
+    return fill_diagonal_offset_(a, val, offset)
\ No newline at end of file

theano/tensor/tests/test_extra_ops.py

@@ -464,3 +464,47 @@ class TestFillDiagonal(utt.InferShapeTester):
                                  numpy.random.rand()],
                                 self.op_class,
                                 warn=False)
+
+class TestFillDiagonalOffset(utt.InferShapeTester):
+
+    rng = numpy.random.RandomState(43)
+
+    def setUp(self):
+        super(TestFillDiagonalOffset, self).setUp()
+        self.op_class = FillDiagonalOffset
+        self.op = fill_diagonal_offset
+
+    def test_perform(self):
+        x = tensor.matrix()
+        y = tensor.scalar()
+        z = tensor.scalar()
+
+        f = function([x, y, z], fill_diagonal(x, y, z))
+        for shp in [(8, 8), (5, 8), (8, 5)]:
+            a = numpy.random.rand(*shp).astype(config.floatX)
+            val = numpy.cast[config.floatX](numpy.random.rand())
+            out = f(a, val, offset)
+            # We can't use numpy.fill_diagonal as it is bugged.
+            assert numpy.allclose(numpy.diag(out, offset), val)
+            assert (out == val).sum() == min(a.shape)
+
+    def test_gradient(self):
+        utt.verify_grad(fill_diagonal_offset, [numpy.random.rand(5, 8),
+                                        numpy.random.rand()],
+                        n_tests=1, rng=TestFillDiagonalOffset.rng)
+        utt.verify_grad(fill_diagonal_offset, [numpy.random.rand(8, 5),
+                                        numpy.random.rand()],
+                        n_tests=1, rng=TestFillDiagonalOffset.rng)
+
+    def test_infer_shape(self):
+        x = tensor.dmatrix()
+        y = tensor.dscalar()
+        z = tensor.dscalar()
+        self._compile_and_check([x, y, z], [self.op(x, y, z)],
+                                [numpy.random.rand(8, 5),
+                                 numpy.random.rand(),
+                                 numpy.random.randint(0,5)],
+                                self.op_class)
+
+if __name__ == '__main__':
+    unittest.main()
\ No newline at end of file

theano/tensor/tests/test_fill_diagonal_offset.py

+import numpy as np
+import numpy
+import unittest
+import pdb
+import theano
+from theano.tests import unittest_tools as utt
+
+from theano.tensor.extra_ops import (CumsumOp, cumsum, CumprodOp, cumprod,
+                                     BinCountOp, bincount, DiffOp, diff,
+                                     squeeze, RepeatOp, repeat, Bartlett, bartlett,
+                                     FillDiagonal, fill_diagonal, FillDiagonalOffset, 
+                                     fill_diagonal_offset)
+from theano import tensor as T
+from theano import config, tensor, function
+
+
+numpy_ver = [int(n) for n in numpy.__version__.split('.')[:2]]
+numpy_16 = bool(numpy_ver >= [1, 6])
+
+
+class TestFillDiagonalOffset(utt.InferShapeTester):
+
+    rng = numpy.random.RandomState(43)
+
+    def setUp(self):
+        super(TestFillDiagonalOffset, self).setUp()
+        self.op_class = FillDiagonalOffset
+        self.op = fill_diagonal_offset
+
+    def test_perform(self):
+        x = tensor.matrix()
+        y = tensor.scalar()
+        z = tensor.scalar()
+
+        test_offset = numpy.array(numpy.random.randint(0,5),
+        	dtype = config.floatX)
+
+        f = function([x, y, z], fill_diagonal_offset(x, y, z))
+        for shp in [(8, 8), (5, 8), (8, 5)]:
+            a = numpy.random.rand(*shp).astype(config.floatX)
+            val = numpy.cast[config.floatX](numpy.random.rand())
+            out = f(a, val, test_offset)
+            # We can't use numpy.fill_diagonal as it is bugged.
+            assert numpy.allclose(numpy.diag(out, test_offset), val)
+            pdb.set_trace()
+            assert (out == val).sum() == min(a.shape)
+
+    def test_gradient(self):
+    	test_offset = numpy.array(numpy.random.randint(0,5),
+        				dtype = config.floatX)
+        utt.verify_grad(fill_diagonal_offset, [numpy.random.rand(5, 8),
+                                        numpy.random.rand(),
+                                        test_offset],
+                        n_tests=1, rng=TestFillDiagonalOffset.rng)
+        utt.verify_grad(fill_diagonal_offset, [numpy.random.rand(8, 5),
+                                        numpy.random.rand(),
+                                        test_offset],
+                        n_tests=1, rng=TestFillDiagonalOffset.rng)
+
+    def test_infer_shape(self):
+        x = tensor.dmatrix()
+        y = tensor.dscalar()
+        z = tensor.dscalar()
+        test_offset = numpy.array(numpy.random.randint(0,5),
+        				dtype = config.floatX)
+        self._compile_and_check([x, y, z], [self.op(x, y, z)],
+                                [numpy.random.rand(8, 5),
+                                 numpy.random.rand(),
+                                 test_offset],
+                                self.op_class)
+
+if __name__ == '__main__':
+    unittest.main()
\ No newline at end of file