Allow passing static shape to tensor creation helpers (#118)

* Allow passing static shape to tensor creation helpers

* Also default dtype to "floatX" when using `tensor`

* 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/shape.py

@@ -641,7 +641,7 @@ def make_node(self, x, shp):
                 except NotScalarConstantError:
                     pass
 
-        return Apply(self, [x, shp], [tensor(x.type.dtype, shape=out_shape)])
+        return Apply(self, [x, shp], [tensor(dtype=x.type.dtype, shape=out_shape)])
 
     def perform(self, node, inp, out_, params):
         x, shp = inp