Remove broadcastable shape arguments to TensorType

Author: brandonwillard

aesara/compile/debugmode.py

@@ -848,17 +848,17 @@ def _get_preallocated_maps(
         or "ALL" in prealloc_modes
     ):
         max_ndim = 0
-        rev_out_broadcastable = []
+        rev_out_shape = []
         for r in considered_outputs:
             if isinstance(r.type, TensorType):
                 if max_ndim < r.ndim:
-                    rev_out_broadcastable += [True] * (r.ndim - max_ndim)
+                    rev_out_shape += [1] * (r.ndim - max_ndim)
                     max_ndim = r.ndim
-                assert len(rev_out_broadcastable) == max_ndim
+                assert len(rev_out_shape) == max_ndim
 
-                for i, b in enumerate(r.broadcastable[::-1]):
-                    rev_out_broadcastable[i] = rev_out_broadcastable[i] and b
-        out_broadcastable = rev_out_broadcastable[::-1]
+                for i, s in enumerate(r.type.shape[::-1]):
+                    rev_out_shape[i] = 1 if rev_out_shape[i] == 1 and s == 1 else None
+        out_shape = rev_out_shape[::-1]
 
     if "strided" in prealloc_modes or "ALL" in prealloc_modes:
         check_ndim = config.DebugMode__check_preallocated_output_ndim
@@ -887,14 +887,14 @@ def _get_preallocated_maps(
         # Moreover, to avoid memory problems, we do not test with strides
         # 2 and -2 on those dimensions.
         step_signs_list = []
-        for b in out_broadcastable[-check_ndim:]:
-            if b:
+        for s in out_shape[-check_ndim:]:
+            if s == 1:
                 step_signs_list.append((1,))
             else:
                 step_signs_list.append((-1, 1))
 
         # Use the same step on all dimensions before the last check_ndim.
-        if all(out_broadcastable[:-check_ndim]):
+        if all(s == 1 for s in out_shape[:-check_ndim]):
             step_signs_list = [(1,)] + step_signs_list
         else:
             step_signs_list = [(-1, 1)] + step_signs_list
@@ -905,7 +905,7 @@ def _get_preallocated_maps(
 
                 # First, the dimensions above check_ndim, then the other ones
                 # Do not test with 2 or -2 for dimensions above check_ndim
-                steps = [step_signs[0]] * len(out_broadcastable[:-check_ndim])
+                steps = [step_signs[0]] * len(out_shape[:-check_ndim])
                 steps += [s * step_size for s in step_signs[1:]]
 
                 name = f"strided{tuple(steps)}"
@@ -932,8 +932,8 @@ def _get_preallocated_maps(
 
     if "wrong_size" in prealloc_modes or "ALL" in prealloc_modes:
         # For each dimension, try size-1, size, size+1
-        for dim, b in enumerate(out_broadcastable):
-            if b:
+        for dim, s in enumerate(out_shape):
+            if s == 1:
                 # The shape has to be 1
                 continue
 
@@ -947,11 +947,11 @@ def _get_preallocated_maps(
                 for r in considered_outputs:
                     if isinstance(r.type, TensorType):
                         r_shape_diff = shape_diff[: r.ndim]
-                        out_shape = [
+                        new_buf_shape = [
                             max((s + sd), 0)
                             for s, sd in zip(r_vals[r].shape, r_shape_diff)
                         ]
-                        new_buf = np.empty(out_shape, dtype=r.type.dtype)
+                        new_buf = np.empty(new_buf_shape, dtype=r.type.dtype)
                         new_buf[...] = np.asarray(def_val).astype(r.type.dtype)
                         wrong_size[r] = new_buf
 

aesara/gradient.py

@@ -1802,13 +1802,11 @@ def verify_grad(
         mode=mode,
     )
 
-    tensor_pt = [
-        aesara.tensor.type.TensorType(
-            aesara.tensor.as_tensor_variable(p).dtype,
-            aesara.tensor.as_tensor_variable(p).broadcastable,
-        )(name=f"input {i}")
-        for i, p in enumerate(pt)
-    ]
+    tensor_pt = []
+    for i, p in enumerate(pt):
+        p_t = aesara.tensor.as_tensor_variable(p).type()
+        p_t.name = f"input {i}"
+        tensor_pt.append(p_t)
 
     # fun can be either a function or an actual Op instance
     o_output = fun(*tensor_pt)

aesara/link/c/params_type.py

@@ -29,7 +29,7 @@
 
 .. code-block:: python
 
-    params_type = ParamsType(attr1=TensorType('int32', (False, False)), attr2=ScalarType('float64'))
+    params_type = ParamsType(attr1=TensorType('int32', shape=(None, None)), attr2=ScalarType('float64'))
 
 If your op contains attributes ``attr1`` **and** ``attr2``, the default ``op.get_params()``
 implementation will automatically try to look for it and generate an appropriate Params object.
@@ -324,7 +324,7 @@ class ParamsType(CType):
     `ParamsType` constructor takes key-value args.  Key will be the name of the
     attribute in the struct.  Value is the Aesara type of this attribute,
     ie. an instance of (a subclass of) :class:`CType`
-    (eg. ``TensorType('int64', (False,))``).
+    (eg. ``TensorType('int64', (None,))``).
 
     In a Python code any attribute named ``key`` will be available via::
 

aesara/sandbox/multinomial.py

@@ -44,7 +44,9 @@ def make_node(self, pvals, unis, n=1):
             odtype = pvals.dtype
         else:
             odtype = self.odtype
-        out = at.tensor(dtype=odtype, shape=pvals.type.broadcastable)
+        out = at.tensor(
+            dtype=odtype, shape=tuple(1 if s == 1 else None for s in pvals.type.shape)
+        )
         return Apply(self, [pvals, unis, as_scalar(n)], [out])
 
     def grad(self, ins, outgrads):

aesara/sandbox/rng_mrg.py

@@ -379,20 +379,23 @@ def make_node(self, rstate, size):
         # this op should not be called directly.
         #
         # call through MRG_RandomStream instead.
-        broad = []
+        out_shape = ()
         for i in range(self.output_type.ndim):
-            broad.append(at.extract_constant(size[i]) == 1)
-        output_type = self.output_type.clone(shape=broad)()
+            if at.extract_constant(size[i]) == 1:
+                out_shape += (1,)
+            else:
+                out_shape += (None,)
+        output_var = self.output_type.clone(shape=out_shape)()
         rstate = as_tensor_variable(rstate)
         size = as_tensor_variable(size)
-        return Apply(self, [rstate, size], [rstate.type(), output_type])
+        return Apply(self, [rstate, size], [rstate.type(), output_var])
 
     @classmethod
     def new(cls, rstate, ndim, dtype, size):
         v_size = as_tensor_variable(size)
         if ndim is None:
             ndim = get_vector_length(v_size)
-        op = cls(TensorType(dtype, (False,) * ndim))
+        op = cls(TensorType(dtype, shape=(None,) * ndim))
         return op(rstate, v_size)
 
     def perform(self, node, inp, out, params):

aesara/sparse/basic.py

@@ -592,7 +592,7 @@ def make_node(self, csm):
 
         csm = as_sparse_variable(csm)
         assert csm.format in ("csr", "csc")
-        data = TensorType(dtype=csm.type.dtype, shape=(False,))()
+        data = TensorType(dtype=csm.type.dtype, shape=(None,))()
         return Apply(self, [csm], [data, ivector(), ivector(), ivector()])
 
     def perform(self, node, inputs, out):
@@ -994,7 +994,7 @@ def make_node(self, x):
         return Apply(
             self,
             [x],
-            [TensorType(dtype=x.type.dtype, shape=(False, False))()],
+            [TensorType(dtype=x.type.dtype, shape=(None, None))()],
         )
 
     def perform(self, node, inputs, outputs):
@@ -1753,11 +1753,13 @@ def __init__(self, axis=None, sparse_grad=True):
     def make_node(self, x):
         x = as_sparse_variable(x)
         assert x.format in ("csr", "csc")
-        b = ()
+
         if self.axis is not None:
-            b = (False,)
+            out_shape = (None,)
+        else:
+            out_shape = ()
 
-        z = TensorType(shape=b, dtype=x.dtype)()
+        z = TensorType(dtype=x.dtype, shape=out_shape)()
         return Apply(self, [x], [z])
 
     def perform(self, node, inputs, outputs):
@@ -1872,7 +1874,7 @@ def make_node(self, x):
         """
         x = as_sparse_variable(x)
         assert x.format in ("csr", "csc")
-        return Apply(self, [x], [tensor(shape=(False,), dtype=x.dtype)])
+        return Apply(self, [x], [tensor(dtype=x.dtype, shape=(None,))])
 
     def perform(self, node, inputs, outputs):
         (x,) = inputs
@@ -2138,7 +2140,7 @@ def make_node(self, x, y):
         return Apply(
             self,
             [x, y],
-            [TensorType(dtype=out_dtype, shape=y.type.broadcastable)()],
+            [TensorType(dtype=out_dtype, shape=y.type.shape)()],
         )
 
     def perform(self, node, inputs, outputs):
@@ -2621,7 +2623,7 @@ def make_node(self, x, y):
         x, y = as_sparse_variable(x), at.as_tensor_variable(y)
 
         assert y.type.ndim == 2
-        out = TensorType(dtype="uint8", shape=(False, False))()
+        out = TensorType(dtype="uint8", shape=(None, None))()
         return Apply(self, [x, y], [out])
 
     def perform(self, node, inputs, outputs):
@@ -3462,7 +3464,7 @@ def make_node(self, a, b):
             return Apply(
                 self,
                 [a, b],
-                [tensor(dtype_out, (False, b.type.broadcastable[1]))],
+                [tensor(dtype_out, shape=(None, 1 if b.type.shape[1] == 1 else None))],
             )
 
     def perform(self, node, inputs, outputs):
@@ -3593,7 +3595,7 @@ 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, (False,))]
+            self, [a_indices, a_indptr, b, g_ab], [tensor(g_ab.dtype, shape=(None,))]
         )
 
     def perform(self, node, inputs, outputs):
@@ -3726,7 +3728,9 @@ class StructuredDotGradCSR(COp):
     __props__ = ()
 
     def make_node(self, a_indices, a_indptr, b, g_ab):
-        return Apply(self, [a_indices, a_indptr, b, g_ab], [tensor(b.dtype, (False,))])
+        return Apply(
+            self, [a_indices, a_indptr, b, g_ab], [tensor(b.dtype, shape=(None,))]
+        )
 
     def perform(self, node, inputs, outputs):
         (a_indices, a_indptr, b, g_ab) = inputs
@@ -3967,6 +3971,7 @@ def make_node(self, x, y):
             x = as_sparse_variable(x)
         if isinstance(y, scipy.sparse.spmatrix):
             y = as_sparse_variable(y)
+
         x_is_sparse_var = _is_sparse_variable(x)
         y_is_sparse_var = _is_sparse_variable(y)
 
@@ -3978,34 +3983,35 @@ def make_node(self, x, y):
             )
 
         if x_is_sparse_var:
-            broadcast_x = (False,) * x.ndim
+            shape_x = (None,) * x.type.ndim
         else:
             x = at.as_tensor_variable(x)
-            broadcast_x = x.type.broadcastable
+            shape_x = x.type.shape
             assert y.format in ("csr", "csc")
             if x.ndim not in (1, 2):
                 raise TypeError(
                     "Input 0 (0-indexed) must have ndim of "
-                    f"1 or 2, {int(x.ndim)} given."
+                    f"1 or 2, {int(x.type.ndim)} given."
                 )
 
         if y_is_sparse_var:
-            broadcast_y = (False,) * y.ndim
+            shape_y = (None,) * y.type.ndim
         else:
             y = at.as_tensor_variable(y)
-            broadcast_y = y.type.broadcastable
+            shape_y = y.type.shape
             assert x.format in ("csr", "csc")
             if y.ndim not in (1, 2):
                 raise TypeError(
                     "Input 1 (1-indexed) must have ndim of "
-                    f"1 or 2, {int(y.ndim)} given."
+                    f"1 or 2, {int(y.type.ndim)} given."
                 )
 
-        if len(broadcast_y) == 2:
-            broadcast_out = broadcast_x[:-1] + broadcast_y[1:]
-        elif len(broadcast_y) == 1:
-            broadcast_out = broadcast_x[:-1]
-        return Apply(self, [x, y], [tensor(dtype=dtype_out, shape=broadcast_out)])
+        if len(shape_y) == 2:
+            shape_out = shape_x[:-1] + shape_y[1:]
+        elif len(shape_y) == 1:
+            shape_out = shape_x[:-1]
+
+        return Apply(self, [x, y], [tensor(dtype=dtype_out, shape=shape_out)])
 
     def perform(self, node, inputs, out):
         x, y = inputs
@@ -4126,21 +4132,21 @@ def make_node(self, alpha, x, y, z):
         alpha = at.as_tensor_variable(alpha)
         z = at.as_tensor_variable(z)
 
-        assert z.ndim == 2
-        assert alpha.type.broadcastable == (True,) * alpha.ndim
+        assert z.type.ndim == 2
+        assert alpha.type.shape == (1,) * alpha.type.ndim
         if not _is_sparse_variable(x):
             x = at.as_tensor_variable(x)
             assert y.format in ("csr", "csc")
-            assert x.ndim == 2
+            assert x.type.ndim == 2
         if not _is_sparse_variable(y):
             y = at.as_tensor_variable(y)
             assert x.format in ("csr", "csc")
-            assert y.ndim == 2
+            assert y.type.ndim == 2
 
         return Apply(
             self,
             [alpha, x, y, z],
-            [tensor(dtype=dtype_out, shape=(False, False))],
+            [tensor(dtype=dtype_out, shape=(None, None))],
         )
 
     def perform(self, node, inputs, outputs):