Move theano.compile.ops tests into their test module

tests/compile/test_ops.py

-"""
- Tests for the Op decorator
-"""
-
-
 import pickle
 
 import numpy as np
+import pytest
 
 import theano
 from tests import unittest_tools as utt
-from theano import function
-from theano.compile import as_op
-from theano.tensor import dmatrix, dvector
+from theano import config, function
+from theano.compile.ops import Rebroadcast, SpecifyShape, as_op
+from theano.tensor.basic import (
+    TensorType,
+    dmatrix,
+    dtensor4,
+    dvector,
+    ivector,
+    matrix,
+    vector,
+)
 
 
-# This is for test_pickle, since the function still has to be
-# reachable from pickle (as in it cannot be defined inline)
 @as_op([dmatrix, dmatrix], dmatrix)
 def mul(a, b):
+    """
+    This is for test_pickle, since the function still has to be
+    reachable from pickle (as in it cannot be defined inline)
+    """
     return a * b
 
 
@@ -87,3 +93,130 @@ class TestOpDecorator(utt.InferShapeTester):
 
 def test_shape_i_hash():
     assert isinstance(theano.tensor.opt.Shape_i(np.int64(1)).__hash__(), int)
+
+
+class TestSpecifyShape(utt.InferShapeTester):
+    mode = None
+    input_type = TensorType
+
+    def shortDescription(self):
+        return None
+
+    def test_bad_shape(self):
+        # Test that at run time we raise an exception when the shape
+        # is not the one specified
+        specify_shape = SpecifyShape()
+
+        x = vector()
+        xval = np.random.rand(2).astype(config.floatX)
+        f = theano.function([x], specify_shape(x, [2]), mode=self.mode)
+        f(xval)
+        xval = np.random.rand(3).astype(config.floatX)
+        with pytest.raises(AssertionError):
+            f(xval)
+
+        assert isinstance(
+            [n for n in f.maker.fgraph.toposort() if isinstance(n.op, SpecifyShape)][0]
+            .inputs[0]
+            .type,
+            self.input_type,
+        )
+
+        x = matrix()
+        xval = np.random.rand(2, 3).astype(config.floatX)
+        f = theano.function([x], specify_shape(x, [2, 3]), mode=self.mode)
+        assert isinstance(
+            [n for n in f.maker.fgraph.toposort() if isinstance(n.op, SpecifyShape)][0]
+            .inputs[0]
+            .type,
+            self.input_type,
+        )
+        f(xval)
+        for shape_ in [(1, 3), (2, 2), (5, 5)]:
+            xval = np.random.rand(*shape_).astype(config.floatX)
+            with pytest.raises(AssertionError):
+                f(xval)
+
+    def test_bad_number_of_shape(self):
+        # Test that the number of dimensions provided is good
+        specify_shape = SpecifyShape()
+
+        x = vector()
+        shape_vec = ivector()
+        xval = np.random.rand(2).astype(config.floatX)
+        with pytest.raises(AssertionError):
+            specify_shape(x, [])
+        with pytest.raises(AssertionError):
+            specify_shape(x, [2, 2])
+
+        f = theano.function([x, shape_vec], specify_shape(x, shape_vec), mode=self.mode)
+        assert isinstance(
+            [n for n in f.maker.fgraph.toposort() if isinstance(n.op, SpecifyShape)][0]
+            .inputs[0]
+            .type,
+            self.input_type,
+        )
+        with pytest.raises(AssertionError):
+            f(xval, [])
+        with pytest.raises(AssertionError):
+            f(xval, [2, 2])
+
+        x = matrix()
+        xval = np.random.rand(2, 3).astype(config.floatX)
+        for shape_ in [(), (1,), (2, 3, 4)]:
+            with pytest.raises(AssertionError):
+                specify_shape(x, shape_)
+            f = theano.function(
+                [x, shape_vec], specify_shape(x, shape_vec), mode=self.mode
+            )
+            assert isinstance(
+                [
+                    n
+                    for n in f.maker.fgraph.toposort()
+                    if isinstance(n.op, SpecifyShape)
+                ][0]
+                .inputs[0]
+                .type,
+                self.input_type,
+            )
+            with pytest.raises(AssertionError):
+                f(xval, shape_)
+
+    def test_infer_shape(self):
+        rng = np.random.RandomState(3453)
+        adtens4 = dtensor4()
+        aivec = ivector()
+        aivec_val = [3, 4, 2, 5]
+        adtens4_val = rng.rand(*aivec_val)
+        self._compile_and_check(
+            [adtens4, aivec],
+            [SpecifyShape()(adtens4, aivec)],
+            [adtens4_val, aivec_val],
+            SpecifyShape,
+        )
+
+
+class TestRebroadcast(utt.InferShapeTester):
+    def test_rebroadcast(self):
+        rng = np.random.RandomState(3453)
+        # Rebroadcast
+        adtens4 = dtensor4()
+        adict = [(0, False), (1, True), (2, False), (3, True)]
+        adtens4_val = rng.rand(2, 1, 3, 1)
+        self._compile_and_check(
+            [adtens4],
+            [Rebroadcast(*adict)(adtens4)],
+            [adtens4_val],
+            Rebroadcast,
+            warn=False,
+        )
+
+        adtens4_bro = TensorType("float64", (True, True, True, False))()
+        bdict = [(0, True), (1, False), (2, False), (3, False)]
+        adtens4_bro_val = rng.rand(1, 1, 1, 3)
+        self._compile_and_check(
+            [adtens4_bro],
+            [Rebroadcast(*bdict)(adtens4_bro)],
+            [adtens4_bro_val],
+            Rebroadcast,
+        )

tests/tensor/test_basic.py

@@ -88,11 +88,9 @@ from theano.tensor import (
     Mean,
     NoneConst,
     PermuteRowElements,
-    Rebroadcast,
     Reshape,
     ScalarFromTensor,
     Shape,
-    SpecifyShape,
     Split,
     Tensor,
     TensorFromScalar,
@@ -6377,94 +6375,6 @@ def test_stacklists():
     assert f(x, x, x, x).shape == (2, 2, 4, 4)
 
 
-class TestSpecifyShape:
-    mode = None
-    input_type = TensorType
-
-    def shortDescription(self):
-        return None
-
-    def test_bad_shape(self):
-        # Test that at run time we raise an exception when the shape
-        # is not the one specified
-        specify_shape = SpecifyShape()
-
-        x = vector()
-        xval = np.random.rand(2).astype(config.floatX)
-        f = theano.function([x], specify_shape(x, [2]), mode=self.mode)
-        f(xval)
-        xval = np.random.rand(3).astype(config.floatX)
-        with pytest.raises(AssertionError):
-            f(xval)
-
-        assert isinstance(
-            [n for n in f.maker.fgraph.toposort() if isinstance(n.op, SpecifyShape)][0]
-            .inputs[0]
-            .type,
-            self.input_type,
-        )
-
-        x = matrix()
-        xval = np.random.rand(2, 3).astype(config.floatX)
-        f = theano.function([x], specify_shape(x, [2, 3]), mode=self.mode)
-        assert isinstance(
-            [n for n in f.maker.fgraph.toposort() if isinstance(n.op, SpecifyShape)][0]
-            .inputs[0]
-            .type,
-            self.input_type,
-        )
-        f(xval)
-        for shape_ in [(1, 3), (2, 2), (5, 5)]:
-            xval = np.random.rand(*shape_).astype(config.floatX)
-            with pytest.raises(AssertionError):
-                f(xval)
-
-    def test_bad_number_of_shape(self):
-        # Test that the number of dimensions provided is good
-        specify_shape = SpecifyShape()
-
-        x = vector()
-        shape_vec = ivector()
-        xval = np.random.rand(2).astype(config.floatX)
-        with pytest.raises(AssertionError):
-            specify_shape(x, [])
-        with pytest.raises(AssertionError):
-            specify_shape(x, [2, 2])
-
-        f = theano.function([x, shape_vec], specify_shape(x, shape_vec), mode=self.mode)
-        assert isinstance(
-            [n for n in f.maker.fgraph.toposort() if isinstance(n.op, SpecifyShape)][0]
-            .inputs[0]
-            .type,
-            self.input_type,
-        )
-        with pytest.raises(AssertionError):
-            f(xval, [])
-        with pytest.raises(AssertionError):
-            f(xval, [2, 2])
-
-        x = matrix()
-        xval = np.random.rand(2, 3).astype(config.floatX)
-        for shape_ in [(), (1,), (2, 3, 4)]:
-            with pytest.raises(AssertionError):
-                specify_shape(x, shape_)
-            f = theano.function(
-                [x, shape_vec], specify_shape(x, shape_vec), mode=self.mode
-            )
-            assert isinstance(
-                [
-                    n
-                    for n in f.maker.fgraph.toposort()
-                    if isinstance(n.op, SpecifyShape)
-                ][0]
-                .inputs[0]
-                .type,
-                self.input_type,
-            )
-            with pytest.raises(AssertionError):
-                f(xval, shape_)
-
-
 class TestInferShape(utt.InferShapeTester):
     def test_infer_shape(self):
 
@@ -6718,28 +6628,6 @@ class TestInferShape(utt.InferShapeTester):
             [aiscal], [TensorFromScalar()(aiscal)], [4.0], TensorFromScalar
         )
 
-        # Rebroadcast
-        adtens4 = dtensor4()
-        adict = [(0, False), (1, True), (2, False), (3, True)]
-        adtens4_val = rand(2, 1, 3, 1)
-        self._compile_and_check(
-            [adtens4],
-            [Rebroadcast(*adict)(adtens4)],
-            [adtens4_val],
-            Rebroadcast,
-            warn=False,
-        )
-
-        adtens4_bro = TensorType("float64", (True, True, True, False))()
-        bdict = [(0, True), (1, False), (2, False), (3, False)]
-        adtens4_bro_val = rand(1, 1, 1, 3)
-        self._compile_and_check(
-            [adtens4_bro],
-            [Rebroadcast(*bdict)(adtens4_bro)],
-            [adtens4_bro_val],
-            Rebroadcast,
-        )
-
         # Alloc
         randint = np.random.randint
         adscal = dscalar()
@@ -6819,16 +6707,6 @@ class TestInferShape(utt.InferShapeTester):
             ARange,
         )
 
-        # SpecifyShape
-        aivec_val = [3, 4, 2, 5]
-        adtens4_val = rand(*aivec_val)
-        self._compile_and_check(
-            [adtens4, aivec],
-            [SpecifyShape()(adtens4, aivec)],
-            [adtens4_val, aivec_val],
-            SpecifyShape,
-        )
-
         # Mean
         adtens3_val = rand(3, 4, 5)
         aiscal_val = 2