Move optimizations and tests.

theano/sparse/opt.py

@@ -7,9 +7,9 @@ from theano.sparse import (CSC, CSR, csm_properties, Remove0,
                            register_specialize,
                            csm_grad, csm_grad_c,
                            usmm_csc_dense, usmm)
+from theano.sparse import basic as sparse
 
-from basic import (_structured_dot, _dot,
-                   _is_sparse_variable, usmm_csc_dense_inplace)
+from basic import _is_sparse_variable
 
 
 # This is tested in tests/test_basic.py:UsmmTests
@@ -18,7 +18,7 @@ local_usmm = gof.opt.PatternSub(
      (theano.tensor.mul,
       {'pattern': 'alpha',
        'constraint': lambda expr: numpy.all(expr.type.broadcastable)},
-    (_dot, 'x', 'y'))),
+    (sparse._dot, 'x', 'y'))),
     (usmm, (theano.tensor.neg, 'alpha'), 'x', 'y', 'z'))
 register_specialize(local_usmm, name="local_usmm")
 
@@ -48,7 +48,7 @@ def local_csm_properties_csm(node):
             return ret_var
 
     return False
-register_specialize(local_csm_properties_csm)
+sparse.register_specialize(local_csm_properties_csm)
 
 
 # This is tested in tests/test_basic.py:test_remove0
@@ -57,7 +57,7 @@ def local_inplace_remove0(node):
     """
     Optimization to insert inplace versions of Remove0.
     """
-    if isinstance(node.op, Remove0) and not node.op.inplace:
+    if isinstance(node.op, sparse.Remove0) and not node.op.inplace:
         new_op = node.op.__class__(inplace=True)
         new_node = new_op(*node.inputs)
         return [new_node]
@@ -70,9 +70,9 @@ theano.compile.optdb.register('local_inplace_remove0',
 
 # register a specialization to replace StructuredDot -> StructuredDotCSx
 # This is tested in tests/test_basic.py:792
-@gof.local_optimizer([_structured_dot])
+@gof.local_optimizer([sparse._structured_dot])
 def local_structured_dot(node):
-    if node.op == _structured_dot:
+    if node.op == sparse._structured_dot:
         a, b = node.inputs
         if a.type.format == 'csc':
             a_val, a_ind, a_ptr, a_shape = csm_properties(a)
@@ -95,10 +95,25 @@ def local_structured_dot(node):
 @gof.local_optimizer([usmm_csc_dense])
 def local_usmm_csc_dense_inplace(node):
     if node.op == usmm_csc_dense:
-        return [usmm_csc_dense_inplace(*node.inputs)]
+        return [sparse.usmm_csc_dense_inplace(*node.inputs)]
 register_specialize(local_usmm_csc_dense_inplace, 'inplace')
 
 
+@gof.local_optimizer([sparse.csm_properties])
+def local_csm_properties_csm(node):
+    """if we find csm_properties(CSM(*args)), then we can replace that with the
+    *args directly"""
+    if node.op == sparse.csm_properties:
+        csm, = node.inputs
+        if csm.owner and (csm.owner.op == sparse.CSC or
+                          csm.owner.op == sparse.CSR):
+            # csm.owner.inputs could be broadcastable. In that case, we have
+            # to adjust the broadcasting flag here.
+            ret_var = [theano.tensor.patternbroadcast(i, o.broadcastable)
+                    for i, o in izip(csm.owner.inputs, node.outputs)]
+            return ret_var
+
+
 # This is tested in tests/test_basic.py:UsmmTests
 @gof.local_optimizer([usmm])
 def local_usmm_csx(node):
@@ -124,4 +139,125 @@ def local_usmm_csx(node):
                 return [usmm_csc_dense(alpha, x_val, x_ind, x_ptr,
                                        x_nsparse, y, z)]
     return False
-register_specialize(local_usmm_csx)
+sparse.register_specialize(local_usmm_csx)
+
+
+# register a specialization to replace MulSD -> MulSDCSX
+@gof.local_optimizer([sparse.mul_s_d])
+def local_mul_s_d(node):
+    if node.op == sparse.mul_s_d:
+        x, y = node.inputs
+
+        x_is_sparse_variable = _is_sparse_variable(x)
+
+        if x_is_sparse_variable:
+            svar = x
+            dvar = y
+        else:
+            svar = y
+            dvar = x
+
+        if dvar.type.ndim != 2:
+            return False
+        if svar.type.format == 'csc':
+            CSx = sparse.CSC
+            mul_s_d_csx = sparse.mul_s_d_csc
+        elif svar.type.format == 'csr':
+            CSx = sparse.CSR
+            mul_s_d_csx = sparse.mul_s_d_csr
+        else:
+            raise NotImplemented()
+
+        c_data = mul_s_d_csx(sparse.csm_data(svar),
+                             sparse.csm_indices(svar),
+                             sparse.csm_indptr(svar), dvar)
+
+        return [CSx(c_data,
+                    sparse.csm_indices(svar),
+                    sparse.csm_indptr(svar),
+                    sparse.csm_shape(svar))]
+
+    return False
+sparse.register_specialize(local_mul_s_d)
+
+
+@gof.local_optimizer([sparse.mul_s_v])
+def local_mul_s_v(node):
+    if node.op == sparse.mul_s_v:
+        x, y = node.inputs
+
+        x_is_sparse_variable = _is_sparse_variable(x)
+
+        if x_is_sparse_variable:
+            svar = x
+            dvar = y
+        else:
+            svar = y
+            dvar = x
+
+        if dvar.type.ndim != 1:
+            return False
+        elif svar.type.format == 'csr':
+            CSx = sparse.CSR
+            mul_s_v_csx = sparse.mul_s_v_csr
+        else:
+            return False
+
+        s_val, s_ind, s_ptr, s_shape = sparse.csm_properties(svar)
+
+        c_data = mul_s_v_csx(s_val, s_ind, s_ptr, dvar)
+
+        return [CSx(c_data, s_ind, s_ptr, s_shape)]
+
+    return False
+sparse.register_specialize(local_mul_s_v)
+
+
+@gof.local_optimizer([sparse.structured_add_s_v])
+def local_structured_add_s_v(node):
+    if node.op == sparse.structured_add_s_v:
+        x, y = node.inputs
+
+        x_is_sparse_variable = _is_sparse_variable(x)
+        #y_is_sparse_variable = _is_sparse_variable(y)
+
+        if x_is_sparse_variable:
+            svar = x
+            dvar = y
+        else:
+            svar = y
+            dvar = x
+
+        if dvar.type.ndim != 1:
+            return False
+        elif svar.type.format == 'csr':
+            CSx = sparse.CSR
+            structured_add_s_v_csx = sparse.structured_add_s_v_csr
+        else:
+            return False
+
+        s_val, s_ind, s_ptr, s_shape = sparse.csm_properties(svar)
+
+        c_data = structured_add_s_v_csx(s_val, s_ind, s_ptr, dvar)
+
+        return [CSx(c_data, s_ind, s_ptr, s_shape)]
+
+    return False
+sparse.register_specialize(local_structured_add_s_v)
+
+
+# register a specialization to replace SamplingDot -> SamplingDotCsr
+@gof.local_optimizer([sparse.sampling_dot])
+def local_sampling_dot_csr(node):
+    if node.op == sparse.sampling_dot:
+        x, y, p = node.inputs
+        if p.type.format == 'csr':
+            p_data, p_ind, p_ptr, p_shape = sparse.csm_properties(p)
+
+            z_data, z_ind, z_ptr = sparse.sampling_dot_csr(x, y, p_data,
+                p_ind, p_ptr, p_shape[1])
+
+            return [sparse.CSR(z_data, z_ind, z_ptr, p_shape)]
+    return False
+sparse.register_specialize(local_sampling_dot_csr,
+                           name='local_sampling_dot_csr')

theano/sparse/sandbox/sp2.py

@@ -12,8 +12,8 @@ from theano.sparse.basic import (
     _is_sparse_variable, _is_dense_variable, CSC, CSR,
     csm_properties, csm_data, csm_indices, csm_indptr, csm_shape,
     _is_sparse,
-    Remove0, remove0,
     # To maintain compatibility
+    Remove0, remove0,
     Cast, bcast, wcast, icast, lcast, fcast, dcast, ccast, zcast,
     HStack, hstack, VStack, vstack,
     AddSSData, add_s_s_data,
@@ -27,48 +27,16 @@ from theano.sparse.basic import (
     StructuredAddSVCSR, structured_add_s_v_csr,
     SamplingDot, sampling_dot, SamplingDotCSR, sampling_dot_csr)
 
+# Also for compatibility
+from theano.sparse.opt import (
+    local_mul_s_d, local_mul_s_v,
+    local_structured_add_s_v, local_sampling_dot_csr)
+
 # Alias to maintain compatibility
 EliminateZeros = Remove0
 eliminate_zeros = remove0
 
 
-# register a specialization to replace MulSD -> MulSDCSX
-@gof.local_optimizer([mul_s_d])
-def local_mul_s_d(node):
-    if node.op == mul_s_d:
-        x, y = node.inputs
-
-        x_is_sparse_variable = _is_sparse_variable(x)
-        # y_is_sparse_variable = _is_sparse_variable(y)
-
-        if x_is_sparse_variable:
-            svar = x
-            dvar = y
-        else:
-            svar = y
-            dvar = x
-
-        if dvar.type.ndim != 2:
-            return False
-        if svar.type.format == 'csc':
-            CSx = CSC
-            mul_s_d_csx = mul_s_d_csc
-        elif svar.type.format == 'csr':
-            CSx = CSR
-            mul_s_d_csx = mul_s_d_csr
-        else:
-            raise NotImplemented()
-
-        c_data = mul_s_d_csx(csm_data(svar), csm_indices(svar),
-                             csm_indptr(svar), dvar)
-
-        return [CSx(c_data, csm_indices(svar), csm_indptr(svar),
-                    csm_shape(svar))]
-
-    return False
-register_specialize(local_mul_s_d)
-
-
 class Binomial(gof.op.Op):
     # TODO This op is not an equivalent of numpy.random.binomial. In
     # facts, this does not follow a binomial distribution at all.
@@ -124,88 +92,8 @@ class Binomial(gof.op.Op):
 
     def __str__(self):
         return self.__class__.__name__
+
 csr_fbinomial = Binomial('csr', 'float32')
 csc_fbinomial = Binomial('csc', 'float32')
 csr_dbinomial = Binomial('csr', 'float64')
 csc_dbinomial = Binomial('csc', 'float64')
-
-
-@gof.local_optimizer([mul_s_v])
-def local_mul_s_v(node):
-    if node.op == mul_s_v:
-        x, y = node.inputs
-
-        x_is_sparse_variable = _is_sparse_variable(x)
-
-        if x_is_sparse_variable:
-            svar = x
-            dvar = y
-        else:
-            svar = y
-            dvar = x
-
-        if dvar.type.ndim != 1:
-            return False
-        elif svar.type.format == 'csr':
-            CSx = CSR
-            mul_s_v_csx = mul_s_v_csr
-        else:
-            return False
-
-        s_val, s_ind, s_ptr, s_shape = csm_properties(svar)
-
-        c_data = mul_s_v_csx(s_val, s_ind, s_ptr, dvar)
-
-        return [CSx(c_data, s_ind, s_ptr, s_shape)]
-
-    return False
-register_specialize(local_mul_s_v)
-
-
-@gof.local_optimizer([structured_add_s_v])
-def local_structured_add_s_v(node):
-    if node.op == structured_add_s_v:
-        x, y = node.inputs
-
-        x_is_sparse_variable = _is_sparse_variable(x)
-        #y_is_sparse_variable = _is_sparse_variable(y)
-
-        if x_is_sparse_variable:
-            svar = x
-            dvar = y
-        else:
-            svar = y
-            dvar = x
-
-        if dvar.type.ndim != 1:
-            return False
-        elif svar.type.format == 'csr':
-            CSx = CSR
-            structured_add_s_v_csx = structured_add_s_v_csr
-        else:
-            return False
-
-        s_val, s_ind, s_ptr, s_shape = csm_properties(svar)
-
-        c_data = structured_add_s_v_csx(s_val, s_ind, s_ptr, dvar)
-
-        return [CSx(c_data, s_ind, s_ptr, s_shape)]
-
-    return False
-register_specialize(local_structured_add_s_v)
-
-
-# register a specialization to replace SamplingDot -> SamplingDotCsr
-@gof.local_optimizer([sampling_dot])
-def local_sampling_dot_csr(node):
-    if node.op == sampling_dot:
-        x, y, p = node.inputs
-        if p.type.format == 'csr':
-            p_data, p_ind, p_ptr, p_shape = csm_properties(p)
-
-            z_data, z_ind, z_ptr = sampling_dot_csr(x, y, p_data,
-                p_ind, p_ptr, p_shape[1])
-
-            return [CSR(z_data, z_ind, z_ptr, p_shape)]
-    return False
-register_specialize(local_sampling_dot_csr, name='local_sampling_dot_csr')

theano/sparse/tests/test_basic.py

@@ -941,8 +941,8 @@ class test_structureddot(unittest.TestCase):
 
             theano_time = t1 - t0
             scipy_time = t2 - t1
-            #print 'theano took', theano_time,
-            #print 'scipy took', scipy_time
+            # print 'theano took', theano_time,
+            # print 'scipy took', scipy_time
             overhead_tol = 0.002  # seconds
             overhead_rtol = 1.1  # times as long
             self.assertTrue(numpy.allclose(theano_result, scipy_result))

theano/sparse/tests/test_opt.py

@@ -6,14 +6,12 @@ except ImportError:
     pass  # The variable enable_sparse will be used to disable the test file.
 
 import theano
-from theano import config, tensor
+from theano import sparse, config, tensor
 from theano.sparse import enable_sparse
 from theano.gof.python25 import any
 if not enable_sparse:
     raise SkipTest('Optional package sparse disabled')
 
-from theano.sparse import (CSM, CSMProperties, csm_properties, CSC, CSR,
-                           DenseFromSparse, CSMGrad)
 from theano.sparse.tests.test_basic import random_lil
 
 
@@ -23,13 +21,15 @@ def test_local_csm_properties_csm():
                               tensor.ivector())
     mode = theano.compile.mode.get_default_mode()
     mode = mode.including("specialize", "local_csm_properties_csm")
-    for CS, cast in [(CSC, sp.csc_matrix), (CSR, sp.csr_matrix)]:
+    for CS, cast in [(sparse.CSC, sp.csc_matrix),
+                     (sparse.CSR, sp.csr_matrix)]:
         f = theano.function([data, indices, indptr, shape],
-                            csm_properties(CS(data, indices, indptr, shape)),
+                            sparse.csm_properties(
+                                CS(data, indices, indptr, shape)),
                             mode=mode)
-        #theano.printing.debugprint(f)
-        assert not any(isinstance(node.op, (CSM, CSMProperties)) for node
-                       in f.maker.env.toposort())
+        assert not any(
+            isinstance(node.op, (sparse.CSM, sparse.CSMProperties))
+            for node in f.maker.env.toposort())
         v = cast(random_lil((10, 40),
                             config.floatX, 3))
         f(v.data, v.indices, v.indptr, v.shape)
@@ -45,14 +45,79 @@ def test_local_csm_grad_c():
         mode = theano.compile.Mode(linker='c|py', optimizer='fast_compile')
 
     mode = mode.including("specialize", "local_csm_grad_c")
-    for CS, cast in [(CSC, sp.csc_matrix), (CSR, sp.csr_matrix)]:
-        cost = tensor.sum(DenseFromSparse()(CS(data, indices, indptr, shape)))
+    for CS, cast in [(sparse.CSC, sp.csc_matrix), (sparse.CSR, sp.csr_matrix)]:
+        cost = tensor.sum(sparse.DenseFromSparse()(CS(data, indices, indptr, shape)))
         f = theano.function(
             [data, indices, indptr, shape],
             tensor.grad(cost, data),
             mode=mode)
-        assert not any(isinstance(node.op, CSMGrad) for node
+        assert not any(isinstance(node.op, sparse.CSMGrad) for node
                        in f.maker.env.toposort())
         v = cast(random_lil((10, 40),
                             config.floatX, 3))
         f(v.data, v.indices, v.indptr, v.shape)
+
+
+def test_local_mul_s_d():
+    mode = theano.compile.mode.get_default_mode()
+    mode = mode.including("specialize", "local_mul_s_d")
+
+    for sp_format in sparse.sparse_formats:
+        inputs = [getattr(theano.sparse, sp_format + '_matrix')(),
+                  tensor.matrix()]
+
+        f = theano.function(inputs,
+                            sparse.mul_s_d(*inputs),
+                            mode=mode)
+
+        assert not any(isinstance(node.op, sparse.MulSD) for node
+                       in f.maker.env.toposort())
+
+
+def test_local_mul_s_v():
+    mode = theano.compile.mode.get_default_mode()
+    mode = mode.including("specialize", "local_mul_s_v")
+
+    for sp_format in ['csr']:  # Not implemented for other format
+        inputs = [getattr(theano.sparse, sp_format + '_matrix')(),
+                  tensor.vector()]
+
+        f = theano.function(inputs,
+                            sparse.mul_s_v(*inputs),
+                            mode=mode)
+
+        assert not any(isinstance(node.op, sparse.MulSV) for node
+                       in f.maker.env.toposort())
+
+
+def test_local_structured_add_s_v():
+    mode = theano.compile.mode.get_default_mode()
+    mode = mode.including("specialize", "local_structured_add_s_v")
+
+    for sp_format in ['csr']:  # Not implemented for other format
+        inputs = [getattr(theano.sparse, sp_format + '_matrix')(),
+                  tensor.vector()]
+
+        f = theano.function(inputs,
+                            sparse.structured_add_s_v(*inputs),
+                            mode=mode)
+
+        assert not any(isinstance(node.op, sparse.StructuredAddSV) for node
+                       in f.maker.env.toposort())
+
+
+def test_local_sampling_dot_csr():
+    mode = theano.compile.mode.get_default_mode()
+    mode = mode.including("specialize", "local_sampling_dot_csr")
+
+    for sp_format in ['csr']:  # Not implemented for other format
+        inputs = [tensor.matrix(),
+                  tensor.matrix(),
+                  getattr(theano.sparse, sp_format + '_matrix')()]
+
+        f = theano.function(inputs,
+                            sparse.sampling_dot(*inputs),
+                            mode=mode)
+
+        assert not any(isinstance(node.op, sparse.SamplingDot) for node
+                       in f.maker.env.toposort())

theano/sparse/tests/test_sp2.py

@@ -21,40 +21,6 @@ from theano.sparse.sandbox import sp2 as S2
 from theano.tests import unittest_tools as utt
 from theano.sparse.basic import verify_grad_sparse
 
-# Already in test_basic.py
-# Not used here
-# def as_sparse_format(data, format):
-#     if format == 'csc':
-#         return scipy.sparse.csc_matrix(data)
-#     elif format == 'csr':
-#         return scipy.sparse.csr_matrix(data)
-#     else:
-#         raise NotImplementedError()
-
-
-# Already in test_basic.py
-# Not used here
-# def eval_outputs(outputs):
-#     return compile.function([], outputs)()[0]
-
-
-# Already in test_basic.py
-# Not used here
-# def random_lil(shape, dtype, nnz):
-#     rval = sp.lil_matrix(shape, dtype=dtype)
-#     huge = 2 ** 30
-#     for k in range(nnz):
-#         # set non-zeros in random locations (row x, col y)
-#         idx = np.random.random_integers(huge, size=len(shape)) % shape
-#         value = np.random.rand()
-#         #if dtype *int*, value will always be zeros!
-#         if "int" in dtype:
-#             value = int(value * 100)
-#         rval.__setitem__(
-#                 idx,
-#                 value)
-#     return rval
-
 
 class BinomialTester(utt.InferShapeTester):
     n = tensor.scalar()
@@ -96,72 +62,5 @@ class BinomialTester(utt.InferShapeTester):
                     self.op_class)
 
 
-# ##################
-# Optimization tests
-# ##################
-def test_local_mul_s_d():
-    mode = theano.compile.mode.get_default_mode()
-    mode = mode.including("specialize", "local_mul_s_d")
-
-    for sp_format in sparse.sparse_formats:
-        inputs = [getattr(theano.sparse, sp_format + '_matrix')(),
-                  tensor.matrix()]
-
-        f = theano.function(inputs,
-                            sparse.mul_s_d(*inputs),
-                            mode=mode)
-
-        assert not any(isinstance(node.op, sparse.MulSD) for node
-                       in f.maker.env.toposort())
-
-
-def test_local_mul_s_v():
-    mode = theano.compile.mode.get_default_mode()
-    mode = mode.including("specialize", "local_mul_s_v")
-
-    for sp_format in ['csr']:  # Not implemented for other format
-        inputs = [getattr(theano.sparse, sp_format + '_matrix')(),
-                  tensor.vector()]
-
-        f = theano.function(inputs,
-                            S2.mul_s_v(*inputs),
-                            mode=mode)
-
-        assert not any(isinstance(node.op, S2.MulSV) for node
-                       in f.maker.env.toposort())
-
-
-def test_local_structured_add_s_v():
-    mode = theano.compile.mode.get_default_mode()
-    mode = mode.including("specialize", "local_structured_add_s_v")
-
-    for sp_format in ['csr']:  # Not implemented for other format
-        inputs = [getattr(theano.sparse, sp_format + '_matrix')(),
-                  tensor.vector()]
-
-        f = theano.function(inputs,
-                            S2.structured_add_s_v(*inputs),
-                            mode=mode)
-
-        assert not any(isinstance(node.op, S2.StructuredAddSV) for node
-                       in f.maker.env.toposort())
-
-
-def test_local_sampling_dot_csr():
-    mode = theano.compile.mode.get_default_mode()
-    mode = mode.including("specialize", "local_sampling_dot_csr")
-
-    for sp_format in ['csr']:  # Not implemented for other format
-        inputs = [tensor.matrix(),
-                  tensor.matrix(),
-                  getattr(theano.sparse, sp_format + '_matrix')()]
-
-        f = theano.function(inputs,
-                            S2.sampling_dot(*inputs),
-                            mode=mode)
-
-        assert not any(isinstance(node.op, S2.SamplingDot) for node
-                       in f.maker.env.toposort())
-
 if __name__ == '__main__':
     unittest.main()