Use np.shape in Shape/Shape_i Python implementations

tests/compile/test_ops.py

@@ -5,8 +5,9 @@ import pytest
 
 import theano
 from tests import unittest_tools as utt
-from theano import config, function
-from theano.compile.ops import Rebroadcast, SpecifyShape, as_op
+from theano import change_flags, config, function
+from theano.compile.ops import Rebroadcast, SpecifyShape, as_op, shape, shape_i
+from theano.gof.fg import FunctionGraph
 from theano.tensor.basic import (
     TensorType,
     dmatrix,
@@ -14,8 +15,13 @@ from theano.tensor.basic import (
     dvector,
     ivector,
     matrix,
+    tensor3,
     vector,
 )
+from theano.tensor.opt import ShapeFeature
+from theano.tensor.subtensor import Subtensor
+from theano.tensor.type_other import NoneConst
+from theano.typed_list import make_list
 
 
 @as_op([dmatrix, dmatrix], dmatrix)
@@ -202,7 +208,7 @@ class TestRebroadcast(utt.InferShapeTester):
         # Rebroadcast
         adtens4 = dtensor4()
         adict = [(0, False), (1, True), (2, False), (3, True)]
-        adtens4_val = rng.rand(2, 1, 3, 1)
+        adtens4_val = rng.rand(2, 1, 3, 1).astype(config.floatX)
         self._compile_and_check(
             [adtens4],
             [Rebroadcast(*adict)(adtens4)],
@@ -213,10 +219,35 @@ class TestRebroadcast(utt.InferShapeTester):
 
         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)
+        adtens4_bro_val = rng.rand(1, 1, 1, 3).astype(config.floatX)
         self._compile_and_check(
             [adtens4_bro],
             [Rebroadcast(*bdict)(adtens4_bro)],
             [adtens4_bro_val],
             Rebroadcast,
         )
+
+
+@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(theano.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(), [])

theano/compile/ops.py

@@ -260,7 +260,7 @@ class Shape(Op):
     def perform(self, node, inp, out_):
         (x,) = inp
         (out,) = out_
-        out[0] = theano._asarray(x.shape, dtype="int64")
+        out[0] = theano._asarray(np.shape(x), dtype="int64")
 
     def infer_shape(self, node, in_shapes):
         return [[len(in_shapes[0])]]
@@ -370,9 +370,6 @@ class Shape_i(Op):
         return "%s{%i}" % (self.__class__.__name__, self.i)
 
     def make_node(self, x):
-        # x could be one of a number of types
-        # the only thing we require is that the variable have a .ndim,
-        # and that the value have a .shape
         if not isinstance(x, theano.Variable):
             raise TypeError("x must be Variable with ndim attribute", x)
         if x.ndim <= self.i:
@@ -383,9 +380,9 @@ class Shape_i(Op):
         (x,) = inp
         (out,) = out_
         if out[0] is None:
-            out[0] = theano._asarray(x.shape[self.i], dtype="int64")
+            out[0] = theano._asarray(np.shape(x)[self.i], dtype="int64")
         else:
-            out[0][...] = x.shape[self.i]
+            out[0][...] = np.shape(x)[self.i]
 
     def c_code_cache_version(self):
         version = []
@@ -489,7 +486,7 @@ def shape_i(var, i, fgraph=None):
     # If we are not able to use the shape feature, we should not put
     # Shape_i in the graph. Otherwise, the shape feature optimization
     # won't get applied.
-    return var.shape[i]
+    return shape(var)[i]
 
 
 def shape_i_op(i):