Merge pull request #73 from delallea/tensor_size

theano/sandbox/cuda/tests/test_basic_ops.py

-import sys, time
+import sys, time, unittest
 
 from theano.compile.pfunc import pfunc
 from theano import tensor
@@ -919,6 +919,34 @@ def test_shared_cudandarray():
     assert isinstance(a.type, tcn.CudaNdarrayType)
 
 
+class test_size(unittest.TestCase):
+
+    """
+    Ensure the `size` attribute of CUDA tensors behaves as in numpy.
+    """
+
+    def test_matrix(self):
+        x = cuda.fmatrix()
+        y = numpy.zeros((5, 7), dtype='float32')
+        assert y.size == theano.function([x], x.size)(y)
+
+    def test_vector(self):
+        x = cuda.fvector()
+        y = numpy.zeros(7, dtype='float32')
+        assert y.size == theano.function([x], x.size)(y)
+
+    def test_scalar(self):
+        x = cuda.fscalar()
+        y = numpy.array(7, dtype='float32')
+        assert y.size == theano.function([x], x.size)(y)
+
+    def test_shared(self):
+        # NB: we also test higher order tensors at the same time.
+        y = cuda.CudaNdarray.zeros((1, 2, 3, 4))
+        x = cuda.shared_constructor(y)
+        assert y.size == theano.function([], x.size)()
+
+
 import theano.tensor.tests.test_sharedvar
 #This test the case when the shared constructor view an CudaNdarray as input
 test_shared_options = theano.tensor.tests.test_sharedvar.makeSharedTester(

theano/sparse/basic.py

@@ -163,6 +163,12 @@ class _sparse_py_operators:
     ndim = property(lambda self: self.type.ndim)
     dtype = property(lambda self: self.type.dtype)
 
+    # Note that the `size` attribute of sparse matrices behaves differently
+    # from dense matrices: it is the number of elements stored in the matrix
+    # rather than the total number of elements that may be stored. Note also
+    # that stored zeros *do* count in the size.
+    size = property(lambda self: csm_data(self).size)
+
 
 class SparseVariable(gof.Variable, _sparse_py_operators):
     dtype = property(lambda self: self.type.dtype)

theano/sparse/tests/test_basic.py

@@ -570,6 +570,26 @@ def test_sparse_shared_memory():
     result_ = f_(a,a)
     assert (result_.todense() == result.todense()).all()
 
+
+def test_size():
+    """
+    Ensure the `size` attribute of sparse matrices behaves as in numpy.
+    """
+    for sparse_type in ('csc_matrix', 'csr_matrix'):
+        x = getattr(theano.sparse, sparse_type)()
+        y = getattr(scipy.sparse, sparse_type)((5, 7))
+        get_size = theano.function([x], x.size)
+        def check():
+            assert y.size == get_size(y)
+        # We verify that the size is correctly updated as we store more data
+        # into the sparse matrix (including zeros).
+        check()
+        y[0, 0] = 1
+        check()
+        y[0, 1] = 0
+        check()
+
+
 import theano.tensor.tests.test_sharedvar
 test_shared_options=theano.tensor.tests.test_sharedvar.makeSharedTester(
     shared_constructor_ = theano.sparse.shared,

theano/tensor/basic.py

@@ -1214,6 +1214,7 @@ class _tensor_py_operators:
     T = property(lambda self: transpose(self))
 
     shape = property(lambda self: shape(self))
+    size = property(lambda self: prod(self.shape))
     def reshape(self, shape, ndim=None):
         """Return a reshaped view/copy of this variable.
 

theano/tensor/tests/test_basic.py

@@ -5151,6 +5151,34 @@ class test_complex_mod(unittest.TestCase):
             pass
 
 
+class test_size(unittest.TestCase):
+
+    """
+    Ensure the `size` attribute of tensors behaves as in numpy.
+    """
+
+    def test_matrix(self):
+        x = tensor.matrix()
+        y = numpy.zeros((5, 7))
+        assert y.size == function([x], x.size)(y)
+
+    def test_vector(self):
+        x = tensor.vector()
+        y = numpy.zeros(7)
+        assert y.size == function([x], x.size)(y)
+
+    def test_scalar(self):
+        x = tensor.scalar()
+        y = numpy.array(7)
+        assert y.size == function([x], x.size)(y)
+
+    def test_shared(self):
+        # NB: we also test higher order tensors at the same time.
+        y = numpy.zeros((1, 2, 3, 4))
+        x = tensor.shared(y)
+        assert y.size == function([], x.size)()
+
+
 if __name__ == '__main__':
     if 1:
         unittest.main()