Move iota tests to test_basic.py

Author: Nimish-4

tests/tensor/test_basic.py

@@ -58,6 +58,7 @@
     identity_like,
     infer_static_shape,
     inverse_permutation,
+    iota,
     join,
     make_vector,
     mgrid,
@@ -980,6 +981,29 @@ def test_static_output_type(self):
         assert eye(1, l, 3).type.shape == (1, None)
 
 
+def test_iota():
+    mode = Mode(linker="py", optimizer=None)
+    np.testing.assert_allclose(
+        iota((4, 8), 0).eval(mode=mode),
+        [
+            [0, 0, 0, 0, 0, 0, 0, 0],
+            [1, 1, 1, 1, 1, 1, 1, 1],
+            [2, 2, 2, 2, 2, 2, 2, 2],
+            [3, 3, 3, 3, 3, 3, 3, 3],
+        ],
+    )
+
+    np.testing.assert_allclose(
+        iota((4, 8), 1).eval(mode=mode),
+        [
+            [0, 1, 2, 3, 4, 5, 6, 7],
+            [0, 1, 2, 3, 4, 5, 6, 7],
+            [0, 1, 2, 3, 4, 5, 6, 7],
+            [0, 1, 2, 3, 4, 5, 6, 7],
+        ],
+    )
+
+
 class TestTriangle:
     def test_tri(self):
         def check(dtype, N, M_=None, k=0):

tests/tensor/test_einsum.py

@@ -10,7 +10,7 @@
 from pytensor.graph.op import HasInnerGraph
 from pytensor.tensor.basic import moveaxis
 from pytensor.tensor.blockwise import Blockwise
-from pytensor.tensor.einsum import _delta, _general_dot, _iota, einsum
+from pytensor.tensor.einsum import _delta, _general_dot, einsum
 from pytensor.tensor.shape import Reshape
 from pytensor.tensor.type import tensor
 
@@ -38,29 +38,6 @@ def assert_no_blockwise_in_graph(fgraph: FunctionGraph, core_op=None) -> None:
             assert_no_blockwise_in_graph(inner_fgraph, core_op=core_op)
 
 
-def test_iota():
-    mode = Mode(linker="py", optimizer=None)
-    np.testing.assert_allclose(
-        _iota((4, 8), 0).eval(mode=mode),
-        [
-            [0, 0, 0, 0, 0, 0, 0, 0],
-            [1, 1, 1, 1, 1, 1, 1, 1],
-            [2, 2, 2, 2, 2, 2, 2, 2],
-            [3, 3, 3, 3, 3, 3, 3, 3],
-        ],
-    )
-
-    np.testing.assert_allclose(
-        _iota((4, 8), 1).eval(mode=mode),
-        [
-            [0, 1, 2, 3, 4, 5, 6, 7],
-            [0, 1, 2, 3, 4, 5, 6, 7],
-            [0, 1, 2, 3, 4, 5, 6, 7],
-            [0, 1, 2, 3, 4, 5, 6, 7],
-        ],
-    )
-
-
 def test_delta():
     mode = Mode(linker="py", optimizer=None)
     np.testing.assert_allclose(