Implement _get_vector_length for SpecifyShape

aesara/tensor/shape.py

@@ -498,6 +498,14 @@ class SpecifyShape(COp):
 specify_shape = SpecifyShape()
 
 
+@_get_vector_length.register(SpecifyShape)
+def _get_vector_length_SpecifyShape(op, var):
+    try:
+        return aet.get_scalar_constant_value(var.owner.inputs[1])
+    except NotScalarConstantError:
+        raise ValueError(f"Length of {var} cannot be determined")
+
+
 class Reshape(COp):
     """Perform a reshape operation of the input x to the new shape shp.
     The number of dimensions to which to reshape to (ndim) must be

tests/tensor/test_basic.py

@@ -960,11 +960,6 @@ def test_get_vector_length():
     res = get_vector_length(x)
     assert res == 4
 
-    # Test `Shape`s
-    x = aesara.shared(np.zeros((2, 3, 4, 5)))
-    res = get_vector_length(x.shape)
-    assert res == 4
-
     # Test `Subtensor`s
     x = as_tensor_variable(np.arange(4))
     assert get_vector_length(x[2:4]) == 2

tests/tensor/test_shape.py

@@ -7,6 +7,7 @@ from aesara.compile.ops import DeepCopyOp
 from aesara.configdefaults import config
 from aesara.graph.fg import FunctionGraph
 from aesara.misc.safe_asarray import _asarray
+from aesara.tensor import get_vector_length
 from aesara.tensor.basic import MakeVector, as_tensor_variable, constant
 from aesara.tensor.basic_opt import ShapeFeature
 from aesara.tensor.elemwise import DimShuffle, Elemwise
@@ -509,3 +510,13 @@ def test_shape_i_basics():
 
     with pytest.raises(TypeError):
         Shape_i(0)(scalar())
+
+
+def test_get_vector_length():
+    # Test `Shape`s
+    x = aesara.shared(np.zeros((2, 3, 4, 5)))
+    assert get_vector_length(x.shape) == 4
+
+    # Test `SpecifyShape`
+    x = specify_shape(ivector(), (10,))
+    assert get_vector_length(x) == 10