Make tensor API similar to that of other variable constructors

* Name is now the only optional non-keyword argument for all constructors

Author: ricardoV94

pytensor/sparse/basic.py

@@ -3451,7 +3451,12 @@ def make_node(self, a, b):
             return Apply(
                 self,
                 [a, b],
-                [tensor(dtype_out, shape=(None, 1 if b.type.shape[1] == 1 else None))],
+                [
+                    tensor(
+                        dtype=dtype_out,
+                        shape=(None, 1 if b.type.shape[1] == 1 else None),
+                    )
+                ],
             )
 
     def perform(self, node, inputs, outputs):
@@ -3582,7 +3587,9 @@ class StructuredDotGradCSC(COp):
 
     def make_node(self, a_indices, a_indptr, b, g_ab):
         return Apply(
-            self, [a_indices, a_indptr, b, g_ab], [tensor(g_ab.dtype, shape=(None,))]
+            self,
+            [a_indices, a_indptr, b, g_ab],
+            [tensor(dtype=g_ab.dtype, shape=(None,))],
         )
 
     def perform(self, node, inputs, outputs):
@@ -3716,7 +3723,7 @@ class StructuredDotGradCSR(COp):
 
     def make_node(self, a_indices, a_indptr, b, g_ab):
         return Apply(
-            self, [a_indices, a_indptr, b, g_ab], [tensor(b.dtype, shape=(None,))]
+            self, [a_indices, a_indptr, b, g_ab], [tensor(dtype=b.dtype, shape=(None,))]
         )
 
     def perform(self, node, inputs, outputs):

pytensor/sparse/rewriting.py

@@ -270,7 +270,11 @@ def make_node(self, a_val, a_ind, a_ptr, a_nrows, b):
         r = Apply(
             self,
             [a_val, a_ind, a_ptr, a_nrows, b],
-            [tensor(dtype_out, shape=(None, 1 if b.type.shape[1] == 1 else None))],
+            [
+                tensor(
+                    dtype=dtype_out, shape=(None, 1 if b.type.shape[1] == 1 else None)
+                )
+            ],
         )
         return r
 
@@ -465,7 +469,12 @@ def make_node(self, a_val, a_ind, a_ptr, b):
         r = Apply(
             self,
             [a_val, a_ind, a_ptr, b],
-            [tensor(self.dtype_out, shape=(None, 1 if b.type.shape[1] == 1 else None))],
+            [
+                tensor(
+                    dtype=self.dtype_out,
+                    shape=(None, 1 if b.type.shape[1] == 1 else None),
+                )
+            ],
         )
         return r
 
@@ -705,7 +714,11 @@ def make_node(self, alpha, x_val, x_ind, x_ptr, x_nrows, y, z):
         r = Apply(
             self,
             [alpha, x_val, x_ind, x_ptr, x_nrows, y, z],
-            [tensor(dtype_out, shape=(None, 1 if y.type.shape[1] == 1 else None))],
+            [
+                tensor(
+                    dtype=dtype_out, shape=(None, 1 if y.type.shape[1] == 1 else None)
+                )
+            ],
         )
         return r
 
@@ -1142,7 +1155,9 @@ def make_node(self, a_data, a_indices, a_indptr, b):
         """
         assert b.type.ndim == 2
         return Apply(
-            self, [a_data, a_indices, a_indptr, b], [tensor(b.dtype, shape=(None,))]
+            self,
+            [a_data, a_indices, a_indptr, b],
+            [tensor(dtype=b.dtype, shape=(None,))],
         )
 
     def c_code_cache_version(self):
@@ -1280,7 +1295,9 @@ def make_node(self, a_data, a_indices, a_indptr, b):
         """
         assert b.type.ndim == 2
         return Apply(
-            self, [a_data, a_indices, a_indptr, b], [tensor(b.dtype, shape=(None,))]
+            self,
+            [a_data, a_indices, a_indptr, b],
+            [tensor(dtype=b.dtype, shape=(None,))],
         )
 
     def c_code_cache_version(self):
@@ -1470,7 +1487,9 @@ def make_node(self, a_data, a_indices, a_indptr, b):
         """
         assert b.type.ndim == 1
         return Apply(
-            self, [a_data, a_indices, a_indptr, b], [tensor(b.dtype, shape=(None,))]
+            self,
+            [a_data, a_indices, a_indptr, b],
+            [tensor(dtype=b.dtype, shape=(None,))],
         )
 
     def c_code_cache_version(self):
@@ -1642,7 +1661,9 @@ def make_node(self, a_data, a_indices, a_indptr, b):
         assert a_indptr.type.ndim == 1
         assert b.type.ndim == 1
         return Apply(
-            self, [a_data, a_indices, a_indptr, b], [tensor(b.dtype, shape=(None,))]
+            self,
+            [a_data, a_indices, a_indptr, b],
+            [tensor(dtype=b.dtype, shape=(None,))],
         )
 
     def c_code_cache_version(self):

pytensor/tensor/basic.py

@@ -2882,7 +2882,7 @@ def make_node(self, start, stop, step):
         assert step.ndim == 0
 
         inputs = [start, stop, step]
-        outputs = [tensor(self.dtype, shape=(None,))]
+        outputs = [tensor(dtype=self.dtype, shape=(None,))]
 
         return Apply(self, inputs, outputs)
 

pytensor/tensor/blas.py

@@ -1680,7 +1680,7 @@ def make_node(self, x, y):
             raise TypeError(y)
         if y.type.dtype != x.type.dtype:
             raise TypeError("dtype mismatch to Dot22")
-        outputs = [tensor(x.type.dtype, shape=(x.type.shape[0], y.type.shape[1]))]
+        outputs = [tensor(dtype=x.type.dtype, shape=(x.type.shape[0], y.type.shape[1]))]
         return Apply(self, [x, y], outputs)
 
     def perform(self, node, inp, out):
@@ -1985,7 +1985,7 @@ def make_node(self, x, y, a):
             raise TypeError("Dot22Scalar requires float or complex args", a.dtype)
 
         sz = (x.type.shape[0], y.type.shape[1])
-        outputs = [tensor(x.type.dtype, shape=sz)]
+        outputs = [tensor(dtype=x.type.dtype, shape=sz)]
         return Apply(self, [x, y, a], outputs)
 
     def perform(self, node, inp, out):
@@ -2221,7 +2221,7 @@ def make_node(self, *inputs):
             + inputs[1].type.shape[2:]
         )
         out_shape = tuple(1 if s == 1 else None for s in out_shape)
-        return Apply(self, upcasted_inputs, [tensor(dtype, shape=out_shape)])
+        return Apply(self, upcasted_inputs, [tensor(dtype=dtype, shape=out_shape)])
 
     def perform(self, node, inp, out):
         x, y = inp

pytensor/tensor/io.py

@@ -36,7 +36,7 @@ def __init__(self, dtype, shape, mmap_mode=None):
     def make_node(self, path):
         if isinstance(path, str):
             path = Constant(Generic(), path)
-        return Apply(self, [path], [tensor(self.dtype, shape=self.shape)])
+        return Apply(self, [path], [tensor(dtype=self.dtype, shape=self.shape)])
 
     def perform(self, node, inp, out):
         path = inp[0]
@@ -135,7 +135,7 @@ def make_node(self):
             [],
             [
                 Variable(Generic(), None),
-                tensor(self.dtype, shape=self.static_shape),
+                tensor(dtype=self.dtype, shape=self.static_shape),
             ],
         )
 
@@ -180,7 +180,7 @@ def make_node(self, request, data):
         return Apply(
             self,
             [request, data],
-            [tensor(data.dtype, shape=data.type.shape)],
+            [tensor(dtype=data.dtype, shape=data.type.shape)],
         )
 
     def perform(self, node, inp, out):

pytensor/tensor/math.py

@@ -152,8 +152,8 @@ def make_node(self, x):
             if i not in all_axes
         )
         outputs = [
-            tensor(x.type.dtype, shape=out_shape, name="max"),
-            tensor("int64", shape=out_shape, name="argmax"),
+            tensor(dtype=x.type.dtype, shape=out_shape, name="max"),
+            tensor(dtype="int64", shape=out_shape, name="argmax"),
         ]
         return Apply(self, inputs, outputs)
 
@@ -370,7 +370,7 @@ def make_node(self, x, axis=None):
         # We keep the original broadcastable flags for dimensions on which
         # we do not perform the argmax.
         out_shape = tuple(s for i, s in enumerate(x.type.shape) if i not in all_axes)
-        outputs = [tensor("int64", shape=out_shape, name="argmax")]
+        outputs = [tensor(dtype="int64", shape=out_shape, name="argmax")]
         return Apply(self, inputs, outputs)
 
     def prepare_node(self, node, storage_map, compute_map, impl):
@@ -1922,7 +1922,7 @@ def make_node(self, *inputs):
             sz = sx[:-1]
 
         i_dtypes = [input.type.dtype for input in inputs]
-        outputs = [tensor(aes.upcast(*i_dtypes), shape=sz)]
+        outputs = [tensor(dtype=aes.upcast(*i_dtypes), shape=sz)]
         return Apply(self, inputs, outputs)
 
     def perform(self, node, inp, out):

pytensor/tensor/type.py

@@ -92,6 +92,12 @@ def __init__(
 
         """
 
+        if shape is None and broadcastable is None:
+            raise ValueError("Must pass shape or broadcastable")
+        elif shape is not None and broadcastable is not None:
+            # TOOD: Lift this constraint when we reintroduce broadcastable flags
+            raise ValueError("Cannot pass both shape and broadcastable")
+
         if broadcastable is not None:
             warnings.warn(
                 "The `broadcastable` keyword is deprecated; use `shape`.",
@@ -775,18 +781,30 @@ def values_eq_approx_always_true(a, b):
     version=2,
 )
 
+# Valid static type entries
+ST = Union[int, None]
+
 
 def tensor(
+    name: Optional[str] = None,
+    *,
     dtype: Optional["DTypeLike"] = None,
-    *args,
+    shape: Optional[Tuple[ST, ...]] = None,
     **kwargs,
 ) -> "TensorVariable":
 
+    if name is not None:
+        # Help catching errors with the new tensor API
+        if str(name) == "floatX" or np.obj2sctype(dtype):
+            raise ValueError(
+                f"The first and only positional argument of tensor is now `name`. Got {name}.\n"
+                "This name looks like a dtype, which you should pass as a keyword argument only."
+            )
+
     if dtype is None:
         dtype = config.floatX
 
-    name = kwargs.pop("name", None)
-    return TensorType(dtype, *args, **kwargs)(name=name)
+    return TensorType(dtype=dtype, shape=shape, **kwargs)(name=name)
 
 
 cscalar = TensorType("complex64", ())
@@ -805,6 +823,7 @@ def tensor(
 
 def scalar(
     name: Optional[str] = None,
+    *,
     dtype: Optional["DTypeLike"] = None,
 ) -> "TensorVariable":
     """Return a symbolic scalar variable.
@@ -844,9 +863,6 @@ def scalar(
 lvector = TensorType("int64", shape=(None,))
 
 
-ST = Union[int, None]
-
-
 def _validate_static_shape(shape, ndim: int) -> Tuple[ST, ...]:
 
     if not isinstance(shape, tuple):
@@ -863,6 +879,7 @@ def _validate_static_shape(shape, ndim: int) -> Tuple[ST, ...]:
 
 def vector(
     name: Optional[str] = None,
+    *,
     dtype: Optional["DTypeLike"] = None,
     shape: Optional[Tuple[ST]] = (None,),
 ) -> "TensorVariable":
@@ -908,6 +925,7 @@ def vector(
 
 def matrix(
     name: Optional[str] = None,
+    *,
     dtype: Optional["DTypeLike"] = None,
     shape: Optional[Tuple[ST, ST]] = (None, None),
 ) -> "TensorVariable":
@@ -951,6 +969,7 @@ def matrix(
 
 def row(
     name: Optional[str] = None,
+    *,
     dtype: Optional["DTypeLike"] = None,
     shape: Optional[Tuple[Literal[1], ST]] = (1, None),
 ) -> "TensorVariable":
@@ -994,6 +1013,7 @@ def row(
 
 def col(
     name: Optional[str] = None,
+    *,
     dtype: Optional["DTypeLike"] = None,
     shape: Optional[Tuple[ST, Literal[1]]] = (None, 1),
 ) -> "TensorVariable":
@@ -1035,6 +1055,7 @@ def col(
 
 def tensor3(
     name: Optional[str] = None,
+    *,
     dtype: Optional["DTypeLike"] = None,
     shape: Optional[Tuple[ST, ST, ST]] = (None, None, None),
 ) -> "TensorVariable":
@@ -1074,6 +1095,7 @@ def tensor3(
 
 def tensor4(
     name: Optional[str] = None,
+    *,
     dtype: Optional["DTypeLike"] = None,
     shape: Optional[Tuple[ST, ST, ST, ST]] = (None, None, None, None),
 ) -> "TensorVariable":
@@ -1113,6 +1135,7 @@ def tensor4(
 
 def tensor5(
     name: Optional[str] = None,
+    *,
     dtype: Optional["DTypeLike"] = None,
     shape: Optional[Tuple[ST, ST, ST, ST, ST]] = (None, None, None, None, None),
 ) -> "TensorVariable":
@@ -1152,6 +1175,7 @@ def tensor5(
 
 def tensor6(
     name: Optional[str] = None,
+    *,
     dtype: Optional["DTypeLike"] = None,
     shape: Optional[Tuple[ST, ST, ST, ST, ST, ST]] = (
         None,
@@ -1198,6 +1222,7 @@ def tensor6(
 
 def tensor7(
     name: Optional[str] = None,
+    *,
     dtype: Optional["DTypeLike"] = None,
     shape: Optional[Tuple[ST, ST, ST, ST, ST, ST, ST]] = (
         None,

tests/graph/test_basic.py

@@ -556,13 +556,13 @@ def test_get_var_by_name():
 def test_clone_new_inputs():
     """Make sure that `Apply.clone_with_new_inputs` properly handles `Type` changes."""
 
-    x = at.tensor(np.float64, shape=(None,))
-    y = at.tensor(np.float64, shape=(1,))
+    x = at.tensor(dtype=np.float64, shape=(None,))
+    y = at.tensor(dtype=np.float64, shape=(1,))
 
     z = at.add(x, y)
     assert z.type.shape == (None,)
 
-    x_new = at.tensor(np.float64, shape=(1,))
+    x_new = at.tensor(dtype=np.float64, shape=(1,))
 
     # The output nodes should be reconstructed, because the input types' static
     # shape information increased in specificity

tests/link/numba/test_elemwise.py

@@ -146,7 +146,7 @@ def test_Elemwise(inputs, input_vals, output_fn, exc):
         # `{'drop': [1], 'shuffle': [2, 0], 'augment': [0, 2, 4]}`
         (
             set_test_value(
-                at.tensor(config.floatX, shape=(None, 1, None), name="a"),
+                at.tensor(dtype=config.floatX, shape=(None, 1, None), name="a"),
                 np.array([[[1.0, 2.0]], [[3.0, 4.0]]], dtype=config.floatX),
             ),
             ("x", 2, "x", 0, "x"),
@@ -155,21 +155,21 @@ def test_Elemwise(inputs, input_vals, output_fn, exc):
         # `{'drop': [1], 'shuffle': [0], 'augment': []}`
         (
             set_test_value(
-                at.tensor(config.floatX, shape=(None, 1), name="a"),
+                at.tensor(dtype=config.floatX, shape=(None, 1), name="a"),
                 np.array([[1.0], [2.0], [3.0], [4.0]], dtype=config.floatX),
             ),
             (0,),
         ),
         (
             set_test_value(
-                at.tensor(config.floatX, shape=(None, 1), name="a"),
+                at.tensor(dtype=config.floatX, shape=(None, 1), name="a"),
                 np.array([[1.0], [2.0], [3.0], [4.0]], dtype=config.floatX),
             ),
             (0,),
         ),
         (
             set_test_value(
-                at.tensor(config.floatX, shape=(1, 1, 1), name="a"),
+                at.tensor(dtype=config.floatX, shape=(1, 1, 1), name="a"),
                 np.array([[[1.0]]], dtype=config.floatX),
             ),
             (),