Modify arange to directly determine the shape from the arguments if possible

Author: Abhinav-Khot

pytensor/tensor/basic.py

@@ -3224,13 +3224,29 @@ def __init__(self, dtype):
         self.dtype = dtype
 
     def make_node(self, start, stop, step):
+        types = TensorConstant
+        shape = (None,)
+        # if it is possible to directly determine the shape i.e static shape is present, we find it.
+        if (
+            isinstance(start, types)
+            and isinstance(stop, types)
+            and isinstance(step, types)
+        ):
+            length = max(
+                np.max(np.ceil((stop.value - start.value) / step.value))
+                .astype(int)
+                .item(),
+                0,
+            )
+            shape = (length,)
+
         start, stop, step = map(as_tensor_variable, (start, stop, step))
         assert start.ndim == 0
         assert stop.ndim == 0
         assert step.ndim == 0
 
         inputs = [start, stop, step]
-        outputs = [tensor(dtype=self.dtype, shape=(None,))]
+        outputs = [tensor(dtype=self.dtype, shape=shape)]
 
         return Apply(self, inputs, outputs)
 

tests/tensor/test_basic.py

@@ -2861,6 +2861,12 @@ def test_infer_shape(self, cast_policy):
             assert np.all(f(2) == len(np.arange(0, 2)))
             assert np.all(f(0) == len(np.arange(0, 0)))
 
+    def test_static_shape(self):
+        assert np.arange(1, 10).shape == arange(1, 10).type.shape
+        assert np.arange(10, 1, -1).shape == arange(10, 1, -1).type.shape
+        assert np.arange(1, -9, 2).shape == arange(1, -9, 2).type.shape
+        assert np.arange(1.3, 17.48, 2.67).shape == arange(1.3, 17.48, 2.67).type.shape
+
 
 class TestNdGrid:
     def setup_method(self):