Move shape-related Ops and functions to theano.tensor.shape

Author: brandonwillard

doc/extending/ctype.txt

@@ -565,17 +565,19 @@ default, it will recompile the c code for each process.
 Shape and Shape_i
 =================
 
-We have 2 generic Ops, Shape and Shape_i, that return the shape of any
-Theano Variable that has a shape attribute (Shape_i returns only one of
+We have 2 generic `Op`s, `Shape` and `Shape_i`, that return the shape of any
+Theano `Variable` that has a shape attribute (`Shape_i` returns only one of
 the elements of the shape).
 
 
 .. code-block:: python
 
-   theano.compile.ops.register_shape_c_code(YOUR_TYPE_CLASS, THE_C_CODE, version=())
-   theano.compile.ops.register_shape_i_c_code(YOUR_TYPE_CLASS, THE_C_CODE, CHECK_INPUT, version=())
+   from theano.theano.shape import register_shape_c_code, register_shape_i_c_code
 
-The C code works as the ViewOp. Shape_i has the additional ``i`` parameter
+   register_shape_c_code(YOUR_TYPE_CLASS, THE_C_CODE, version=())
+   register_shape_i_c_code(YOUR_TYPE_CLASS, THE_C_CODE, CHECK_INPUT, version=())
+
+The C code works as the `ViewOp`. `Shape_i` has the additional ``i`` parameter
 that you can use with ``%(i)s``.
 
 In your CHECK_INPUT, you must check that the input has enough dimensions to

doc/library/gradient.txt

@@ -36,7 +36,7 @@ List of Implemented R op
 See the :ref:`gradient tutorial <tutcomputinggrads>` for the R op documentation.
 
 list of ops that support R-op:
- * with test [Most is tensor/tests/test_rop.py]
+ * with test [Most is tests/tensor/test_rop.py]
     * SpecifyShape
     * MaxAndArgmax
     * Subtensor
@@ -52,7 +52,7 @@ list of ops that support R-op:
     * Reshape
     * Flatten
     * DimShuffle
-    * Scan [In scan/tests/test_scan.test_rop]
+    * Scan [In tests/scan/test_basic.test_rop]
 
  * without test
     * Split

doc/tutorial/shape_info.txt

@@ -118,7 +118,7 @@ upgrade.  Here is the current state of what can be done:
 
 >>> import theano
 >>> x = theano.tensor.matrix()
->>> x_specify_shape = theano.compile.ops.specify_shape(x, (2, 2))
+>>> x_specify_shape = theano.tensor.specify_shape(x, (2, 2))
 >>> f = theano.function([x], (x_specify_shape ** 2).shape)
 >>> theano.printing.debugprint(f) # doctest: +NORMALIZE_WHITESPACE
 DeepCopyOp [id A] ''   0

tests/compile/test_ops.py

@@ -1,28 +1,12 @@
 import pickle
 
 import numpy as np
-import pytest
 
-import theano
 from tests import unittest_tools as utt
 from theano import function
-from theano.compile.ops import Rebroadcast, SpecifyShape, as_op, shape, shape_i
+from theano.compile.ops import Rebroadcast, as_op
 from theano.configdefaults import config
-from theano.graph.fg import FunctionGraph
-from theano.tensor.opt import ShapeFeature
-from theano.tensor.subtensor import Subtensor
-from theano.tensor.type import (
-    TensorType,
-    dmatrix,
-    dtensor4,
-    dvector,
-    ivector,
-    matrix,
-    tensor3,
-    vector,
-)
-from theano.tensor.type_other import NoneConst
-from theano.typed_list import make_list
+from theano.tensor.type import TensorType, dmatrix, dtensor4, dvector
 
 
 @as_op([dmatrix, dmatrix], dmatrix)
@@ -98,111 +82,6 @@ def test_pickle(self):
         assert m2.owner.op == m.owner.op
 
 
-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)
@@ -227,28 +106,3 @@ def test_rebroadcast(self):
             [adtens4_bro_val],
             Rebroadcast,
         )
-
-
-@config.change_flags(compute_test_value="raise")
-def test_nonstandard_shapes():
-    a = tensor3(config.floatX)
-    a.tag.test_value = np.random.random((2, 3, 4)).astype(config.floatX)
-    b = tensor3(config.floatX)
-    b.tag.test_value = np.random.random((2, 3, 4)).astype(config.floatX)
-
-    tl = make_list([a, b])
-    tl_shape = shape(tl)
-    assert np.array_equal(tl_shape.get_test_value(), (2, 2, 3, 4))
-
-    # There's no `FunctionGraph`, so it should return a `Subtensor`
-    tl_shape_i = shape_i(tl, 0)
-    assert isinstance(tl_shape_i.owner.op, Subtensor)
-    assert tl_shape_i.get_test_value() == 2
-
-    tl_fg = FunctionGraph([a, b], [tl], features=[ShapeFeature()])
-    tl_shape_i = shape_i(tl, 0, fgraph=tl_fg)
-    assert not isinstance(tl_shape_i.owner.op, Subtensor)
-    assert tl_shape_i.get_test_value() == 2
-
-    none_shape = shape(NoneConst)
-    assert np.array_equal(none_shape.get_test_value(), [])

tests/gpuarray/test_basic_ops.py

@@ -14,7 +14,6 @@
     TestReshape,
 )
 from tests.tensor.utils import rand, safe_make_node
-from theano.compile.ops import Shape_i
 from theano.gpuarray.basic_ops import (
     GpuAlloc,
     GpuAllocEmpty,
@@ -36,6 +35,7 @@
 from theano.gpuarray.type import GpuArrayType, get_context, gpuarray_shared_constructor
 from theano.tensor.basic import Alloc, Split, alloc
 from theano.tensor.opt import MakeVector
+from theano.tensor.shape import Shape, Shape_i
 from theano.tensor.type import TensorType, fmatrix, iscalar, lscalar, matrix
 
 
@@ -352,7 +352,7 @@ def test_shape():
     topo = f.maker.fgraph.toposort()
     assert np.all(f(v) == (3, 4, 5))
     assert len(topo) == 1
-    assert isinstance(topo[0].op, theano.compile.ops.Shape)
+    assert isinstance(topo[0].op, Shape)
 
 
 def test_gpu_contiguous():

tests/gpuarray/test_dnn.py

@@ -44,6 +44,7 @@
 )
 from theano.tensor.nnet.corr import CorrMM
 from theano.tensor.nnet.corr3d import Corr3dMM
+from theano.tensor.shape import reshape
 from theano.tensor.signal.pool import (
     AveragePoolGrad,
     MaxPoolGrad,
@@ -2813,7 +2814,7 @@ def test_dnn_spatialtf():
 
     def spatialtf_cpu(inp, theta, scale_height, scale_width, border_mode="nearest"):
         num_batch, num_channels, height, width = inp.shape
-        theta = tt.reshape(theta, (-1, 2, 3))
+        theta = reshape(theta, (-1, 2, 3))
 
         # grid of (x_t, y_t, 1), eq (1) in ref [1]
         out_height = tt.cast(tt.ceil(height * scale_height), "int64")
@@ -2832,7 +2833,7 @@ def spatialtf_cpu(inp, theta, scale_height, scale_width, border_mode="nearest"):
             input_dim, x_s_flat, y_s_flat, out_height, out_width, border_mode
         )
 
-        output = tt.reshape(
+        output = reshape(
             input_transformed, (num_batch, out_height, out_width, num_channels)
         )
         output = output.dimshuffle(0, 3, 1, 2)  # dimshuffle to conv format

tests/gpuarray/test_opt.py

@@ -55,14 +55,16 @@
 
 
 def _check_stack_trace(thing):
+    from theano.tensor.shape import Shape, Shape_i
+
     def _ops_to_check(op):
         if not isinstance(op, theano.graph.op.Op):
             op = op.op  # assume it is an apply node
         return not isinstance(
             op,
             (
-                theano.compile.ops.Shape_i,
-                theano.compile.ops.Shape,
+                Shape_i,
+                Shape,
                 theano.compile.ops.DeepCopyOp,
                 theano.tensor.opt.MakeVector,
                 theano.tensor.subtensor.Subtensor,

tests/gpuarray/test_type.py

@@ -15,6 +15,7 @@
 from theano.compile import DeepCopyOp, Rebroadcast, ViewOp
 from theano.configdefaults import config
 from theano.gpuarray.type import GpuArrayType, gpuarray_shared_constructor
+from theano.tensor.shape import specify_shape
 from theano.tensor.type import row
 
 
@@ -80,7 +81,7 @@ def test_specify_shape():
     for dtype in ["float16", "float32"]:
         a = rand_gpuarray(20, dtype=dtype)
         g = GpuArrayType(dtype=dtype, broadcastable=(False,))("g")
-        f = theano.function([g], theano.compile.ops.specify_shape(g, [20]))
+        f = theano.function([g], specify_shape(g, [20]))
         f(a)
 
 

tests/link/test_jax.py

@@ -7,7 +7,6 @@
 from theano.graph.fg import FunctionGraph
 from theano.graph.optdb import Query
 from theano.link.jax import JAXLinker
-from theano.tensor import basic as tt_basic
 from theano.tensor import blas as tt_blas
 from theano.tensor import elemwise as tt_elemwise
 from theano.tensor import extra_ops as tt_extra_ops
@@ -16,6 +15,7 @@
 from theano.tensor import opt as tt_opt
 from theano.tensor import slinalg as tt_slinalg
 from theano.tensor import subtensor as tt_subtensor
+from theano.tensor.shape import Shape, Shape_i, SpecifyShape, reshape
 from theano.tensor.type import (
     dscalar,
     dvector,
@@ -139,35 +139,37 @@ def compare_shape_dtype(x, y):
     compare_jax_and_py(x_fg, [np.ones(10, dtype=theano.config.floatX)])
 
 
-def test_jax_compile_ops():
-    x = theano.compile.ops.DeepCopyOp()(tt.as_tensor_variable(1.1))
-    x_fg = FunctionGraph([], [x])
-
-    compare_jax_and_py(x_fg, [])
-
+def test_jax_shape_ops():
     x_np = np.zeros((20, 3))
-    x = theano.compile.ops.Shape()(tt.as_tensor_variable(x_np))
+    x = Shape()(tt.as_tensor_variable(x_np))
     x_fg = FunctionGraph([], [x])
 
     compare_jax_and_py(x_fg, [], must_be_device_array=False)
 
-    x = theano.compile.ops.Shape_i(1)(tt.as_tensor_variable(x_np))
+    x = Shape_i(1)(tt.as_tensor_variable(x_np))
     x_fg = FunctionGraph([], [x])
 
     compare_jax_and_py(x_fg, [], must_be_device_array=False)
 
-    x = theano.compile.ops.SpecifyShape()(tt.as_tensor_variable(x_np), (20, 3))
+    x = SpecifyShape()(tt.as_tensor_variable(x_np), (20, 3))
     x_fg = FunctionGraph([], [x])
 
     compare_jax_and_py(x_fg, [])
 
     with theano.config.change_flags(compute_test_value="off"):
-        x = theano.compile.ops.SpecifyShape()(tt.as_tensor_variable(x_np), (2, 3))
+        x = SpecifyShape()(tt.as_tensor_variable(x_np), (2, 3))
         x_fg = FunctionGraph([], [x])
 
         with pytest.raises(AssertionError):
             compare_jax_and_py(x_fg, [])
 
+
+def test_jax_compile_ops():
+    x = theano.compile.ops.DeepCopyOp()(tt.as_tensor_variable(1.1))
+    x_fg = FunctionGraph([], [x])
+
+    compare_jax_and_py(x_fg, [])
+
     x_np = np.zeros((20, 1, 1))
     x = theano.compile.ops.Rebroadcast((0, False), (1, True), (2, False))(
         tt.as_tensor_variable(x_np)
@@ -650,13 +652,13 @@ def test_jax_MakeVector():
 
 def test_jax_Reshape():
     a = vector("a")
-    x = tt_basic.reshape(a, (2, 2))
+    x = reshape(a, (2, 2))
     x_fg = FunctionGraph([a], [x])
     compare_jax_and_py(x_fg, [np.r_[1.0, 2.0, 3.0, 4.0].astype(theano.config.floatX)])
 
     # Test breaking "omnistaging" changes in JAX.
     # See https://github.com/tensorflow/probability/commit/782d0c64eb774b9aac54a1c8488e4f1f96fbbc68
-    x = tt_basic.reshape(a, (a.shape[0] // 2, a.shape[0] // 2))
+    x = reshape(a, (a.shape[0] // 2, a.shape[0] // 2))
     x_fg = FunctionGraph([a], [x])
     compare_jax_and_py(x_fg, [np.r_[1.0, 2.0, 3.0, 4.0].astype(theano.config.floatX)])
 
@@ -665,7 +667,7 @@ def test_jax_Reshape():
 def test_jax_Reshape_nonconcrete():
     a = vector("a")
     b = iscalar("b")
-    x = tt_basic.reshape(a, (b, b))
+    x = reshape(a, (b, b))
     x_fg = FunctionGraph([a, b], [x])
     compare_jax_and_py(
         x_fg, [np.r_[1.0, 2.0, 3.0, 4.0].astype(theano.config.floatX), 2]