Merge branch 'main' into time

Author: anirudhacharya

aesara/compile/debugmode.py

@@ -807,7 +807,7 @@ def _get_preallocated_maps(
         for r in considered_outputs:
             if isinstance(r.type, TensorType):
                 # Build a C-contiguous buffer
-                new_buf = r.type.value_zeros(r_vals[r].shape)
+                new_buf = np.empty(r_vals[r].shape, dtype=r.type.dtype)
                 assert new_buf.flags["C_CONTIGUOUS"]
                 new_buf[...] = np.asarray(def_val).astype(r.type.dtype)
 
@@ -875,7 +875,8 @@ def _get_preallocated_maps(
                         buf_shape.append(s)
                     else:
                         buf_shape.append(s * 2)
-                new_buf = r.type.value_zeros(buf_shape)
+
+                new_buf = np.empty(buf_shape, dtype=r.type.dtype)
                 new_buf[...] = np.asarray(def_val).astype(r.type.dtype)
                 init_strided[r] = new_buf
 
@@ -950,7 +951,7 @@ def _get_preallocated_maps(
                             max((s + sd), 0)
                             for s, sd in zip(r_vals[r].shape, r_shape_diff)
                         ]
-                        new_buf = r.type.value_zeros(out_shape)
+                        new_buf = np.empty(out_shape, dtype=r.type.dtype)
                         new_buf[...] = np.asarray(def_val).astype(r.type.dtype)
                         wrong_size[r] = new_buf
 

aesara/scan/op.py

@@ -161,8 +161,9 @@ def check_broadcast(v1, v2):
     which may wrongly be interpreted as broadcastable.
 
     """
-    if not hasattr(v1, "broadcastable") and not hasattr(v2, "broadcastable"):
+    if not isinstance(v1.type, TensorType) and not isinstance(v2.type, TensorType):
         return
+
     msg = (
         "The broadcast pattern of the output of scan (%s) is "
         "inconsistent with the one provided in `output_info` "
@@ -173,13 +174,13 @@ def check_broadcast(v1, v2):
         "them consistent, e.g. using aesara.tensor."
         "{unbroadcast, specify_broadcastable}."
     )
-    size = min(len(v1.broadcastable), len(v2.broadcastable))
+    size = min(v1.type.ndim, v2.type.ndim)
     for n, (b1, b2) in enumerate(
-        zip(v1.broadcastable[-size:], v2.broadcastable[-size:])
+        zip(v1.type.broadcastable[-size:], v2.type.broadcastable[-size:])
     ):
         if b1 != b2:
-            a1 = n + size - len(v1.broadcastable) + 1
-            a2 = n + size - len(v2.broadcastable) + 1
+            a1 = n + size - v1.type.ndim + 1
+            a2 = n + size - v2.type.ndim + 1
             raise TypeError(msg % (v1.type, v2.type, a1, b1, b2, a2))
 
 
@@ -200,7 +201,7 @@ def copy_var_format(var, as_var):
         rval = as_var.type.filter_variable(rval)
     else:
         tmp = as_var.type.clone(
-            shape=(tuple(var.broadcastable[:1]) + tuple(as_var.broadcastable))
+            shape=(tuple(var.type.shape[:1]) + tuple(as_var.type.shape))
         )
         rval = tmp.filter_variable(rval)
     return rval
@@ -628,6 +629,7 @@ def validate_inner_graph(self):
                 type_input = self.inner_inputs[inner_iidx].type
                 type_output = self.inner_outputs[inner_oidx].type
                 if (
+                    # TODO: Use the `Type` interface for this
                     type_input.dtype != type_output.dtype
                     or type_input.broadcastable != type_output.broadcastable
                 ):
@@ -805,7 +807,9 @@ def tensorConstructor(shape, dtype):
             # output sequence
             o = outputs[idx]
             self.output_types.append(
-                typeConstructor((False,) + o.type.broadcastable, o.type.dtype)
+                # TODO: What can we actually say about the shape of this
+                # added dimension?
+                typeConstructor((None,) + o.type.shape, o.type.dtype)
             )
 
             idx += len(info.mit_mot_out_slices[jdx])
@@ -816,7 +820,9 @@ def tensorConstructor(shape, dtype):
 
         for o in outputs[idx:end]:
             self.output_types.append(
-                typeConstructor((False,) + o.type.broadcastable, o.type.dtype)
+                # TODO: What can we actually say about the shape of this
+                # added dimension?
+                typeConstructor((None,) + o.type.shape, o.type.dtype)
             )
 
         # shared outputs + possibly the ending condition
@@ -1380,11 +1386,13 @@ def prepare_fgraph(self, fgraph):
                         # the output value, possibly inplace, at the end of the
                         # function execution. Also, since an update is defined,
                         # a default value must also be (this is verified by
-                        # DebugMode). Use an array of size 0 with the correct
-                        # ndim and dtype (use a shape of 1 on broadcastable
-                        # dimensions, and 0 on the others).
-                        default_shape = [1 if _b else 0 for _b in inp.broadcastable]
-                        default_val = inp.type.value_zeros(default_shape)
+                        # DebugMode).
+                        # TODO FIXME: Why do we need a "default value" here?
+                        # This sounds like a serious design issue.
+                        default_shape = tuple(
+                            s if s is not None else 0 for s in inp.type.shape
+                        )
+                        default_val = np.empty(default_shape, dtype=inp.type.dtype)
                         wrapped_inp = In(
                             variable=inp,
                             value=default_val,
@@ -2318,8 +2326,8 @@ def infer_shape(self, fgraph, node, input_shapes):
                     # equivalent (if False). Here, we only need the variable.
                     v_shp_i = validator.check(shp_i)
                     if v_shp_i is None:
-                        if hasattr(r, "broadcastable") and r.broadcastable[i]:
-                            shp.append(1)
+                        if r.type.shape[i] is not None:
+                            shp.append(r.type.shape[i])
                         else:
                             shp.append(Shape_i(i)(r))
                     else:

aesara/sparse/type.py

@@ -224,14 +224,6 @@ def get_size(self, shape_info):
             + (shape_info[2] + shape_info[3]) * np.dtype("int32").itemsize
         )
 
-    def value_zeros(self, shape):
-        matrix_constructor = self.format_cls.get(self.format)
-
-        if matrix_constructor is None:
-            raise ValueError(f"Sparse matrix type {self.format} not found in SciPy")
-
-        return matrix_constructor(shape, dtype=self.dtype)
-
     def __eq__(self, other):
         res = super().__eq__(other)
 

aesara/tensor/basic.py

@@ -10,7 +10,7 @@
 from collections.abc import Sequence
 from functools import partial
 from numbers import Number
-from typing import Optional
+from typing import TYPE_CHECKING, Optional
 from typing import Sequence as TypeSequence
 from typing import Tuple, Union
 from typing import cast as type_cast
@@ -68,6 +68,10 @@
 from aesara.tensor.var import TensorConstant, TensorVariable, get_unique_value
 
 
+if TYPE_CHECKING:
+    from aesara.tensor import TensorLike
+
+
 def __oplist_tag(thing, tag):
     tags = getattr(thing, "__oplist_tags", [])
     tags.append(tag)
@@ -1334,11 +1338,25 @@ def identity_like(x, dtype: Optional[Union[str, np.generic, np.dtype]] = None):
     return eye(_x.shape[0], _x.shape[1], k=0, dtype=dtype)
 
 
-def infer_broadcastable(shape):
-    """Infer the broadcastable dimensions for `shape`.
+def infer_static_shape(
+    shape: Union[Variable, TypeSequence[Union[Variable, int]]]
+) -> Tuple[TypeSequence["TensorLike"], TypeSequence[Optional[int]]]:
+    """Infer the static shapes implied by the potentially symbolic elements in `shape`.
+
+    `shape` will be validated and constant folded.  As a result, this function
+    can be expensive and shouldn't be used unless absolutely necessary.
+
+    It mostly exists as a hold-over from pre-static shape times, when it was
+    required in order to produce correct broadcastable arrays and prevent
+    some graphs from being unusable.  Now, it is no longer strictly required,
+    so don't use it unless you want the same shape graphs to be rewritten
+    multiple times during graph construction.
+
+    Returns
+    -------
+    A validated sequence of symbolic shape values, and a sequence of
+    ``None``/``int`` values that can be used as `TensorType.shape` values.
 
-    `shape` will be validated and constant folded in order to determine
-    which dimensions are broadcastable (i.e. equal to ``1``).
     """
     from aesara.tensor.rewriting.basic import topo_constant_folding
     from aesara.tensor.rewriting.shape import ShapeFeature
@@ -1362,9 +1380,10 @@ def check_type(s):
         clone=True,
     )
     folded_shape = rewrite_graph(shape_fg, custom_rewrite=topo_constant_folding).outputs
-
-    bcast = tuple(getattr(s, "data", s) == 1 for s in folded_shape)
-    return sh, bcast
+    static_shape = tuple(
+        s.data.item() if isinstance(s, Constant) else None for s in folded_shape
+    )
+    return sh, static_shape
 
 
 class Alloc(COp):
@@ -1394,15 +1413,15 @@ class Alloc(COp):
 
     def make_node(self, value, *shape):
         v = as_tensor_variable(value)
-        sh, bcast = infer_broadcastable(shape)
+        sh, static_shape = infer_static_shape(shape)
         if v.ndim > len(sh):
             raise TypeError(
                 "The Alloc value to use has more dimensions"
                 " than the specified dimensions",
                 v.ndim,
                 len(sh),
             )
-        otype = TensorType(dtype=v.dtype, shape=bcast)
+        otype = TensorType(dtype=v.dtype, shape=static_shape)
         return Apply(self, [v] + sh, [otype()])
 
     def perform(self, node, inputs, out_):
@@ -3823,8 +3842,8 @@ def typecode(self):
         return np.dtype(self.dtype).num
 
     def make_node(self, *_shape):
-        _shape, bcast = infer_broadcastable(_shape)
-        otype = TensorType(dtype=self.dtype, shape=bcast)
+        _shape, static_shape = infer_static_shape(_shape)
+        otype = TensorType(dtype=self.dtype, shape=static_shape)
         output = otype()
 
         output.tag.values_eq_approx = values_eq_approx_always_true

aesara/tensor/extra_ops.py

@@ -23,6 +23,7 @@
 from aesara.raise_op import Assert
 from aesara.scalar import int32 as int_t
 from aesara.scalar import upcast
+from aesara.scalar.basic import Composite
 from aesara.tensor import basic as at
 from aesara.tensor import get_vector_length
 from aesara.tensor.exceptions import NotScalarConstantError
@@ -1552,16 +1553,32 @@ def broadcast_shape_iter(
                     # be broadcastable or equal to the one non-broadcastable
                     # constant `const_nt_shape_var`.
                     assert_dim = Assert("Could not broadcast dimensions")
+
+                    scalar_nonconst_nb_shapes = [
+                        at.scalar_from_tensor(s)
+                        if isinstance(s.type, TensorType)
+                        else s
+                        for s in nonconst_nb_shapes
+                    ]
+
+                    dummy_nonconst_nb_shapes = [
+                        aes.get_scalar_type(dtype=v.dtype)()
+                        for v in scalar_nonconst_nb_shapes
+                    ]
                     assert_cond = reduce(
                         aes.and_,
                         (
                             aes.or_(
                                 aes.eq(nbv, one_at), aes.eq(nbv, const_nt_shape_var)
                             )
-                            for nbv in nonconst_nb_shapes
+                            for nbv in dummy_nonconst_nb_shapes
                         ),
                     )
-                    bcast_dim = assert_dim(const_nt_shape_var, assert_cond)
+                    assert_cond_op = Composite(dummy_nonconst_nb_shapes, [assert_cond])
+
+                    bcast_dim = assert_dim(
+                        const_nt_shape_var, assert_cond_op(*scalar_nonconst_nb_shapes)
+                    )
                 else:
                     bcast_dim = const_nt_shape_var
             else:
@@ -1579,21 +1596,37 @@ def broadcast_shape_iter(
                     result_dims.append(maybe_non_bcast_shapes[0])
                     continue
 
+                scalar_maybe_non_bcast_shapes = [
+                    at.scalar_from_tensor(s) if isinstance(s.type, TensorType) else s
+                    for s in maybe_non_bcast_shapes
+                ]
+                dummy_maybe_non_bcast_shapes = [
+                    aes.get_scalar_type(dtype=v.dtype)()
+                    for v in scalar_maybe_non_bcast_shapes
+                ]
                 non_bcast_vec = [
                     aes.switch(aes.eq(nbv, 1), -one_at, nbv)
-                    for nbv in maybe_non_bcast_shapes
+                    for nbv in dummy_maybe_non_bcast_shapes
                 ]
                 dim_max = aes.abs(reduce(aes.scalar_maximum, non_bcast_vec))
+                dim_max_op = Composite(dummy_maybe_non_bcast_shapes, [dim_max])
+
+                dummy_dim_max = dim_max_op(*dummy_maybe_non_bcast_shapes)
 
                 assert_dim = Assert("Could not broadcast dimensions")
                 assert_cond = reduce(
                     aes.and_,
                     (
-                        aes.or_(aes.eq(nbv, -one_at), aes.eq(nbv, dim_max))
+                        aes.or_(aes.eq(nbv, -one_at), aes.eq(nbv, dummy_dim_max))
                         for nbv in non_bcast_vec
                     ),
                 )
-                bcast_dim = assert_dim(dim_max, assert_cond)
+                assert_cond_op = Composite(dummy_maybe_non_bcast_shapes, [assert_cond])
+
+                bcast_dim = assert_dim(
+                    dim_max_op(*scalar_maybe_non_bcast_shapes),
+                    assert_cond_op(*scalar_maybe_non_bcast_shapes),
+                )
 
             result_dims.append(bcast_dim)
 
@@ -1613,9 +1646,9 @@ def __call__(self, a, shape, **kwargs):
     def make_node(self, a, *shape):
         a = at.as_tensor_variable(a)
 
-        shape, bcast = at.infer_broadcastable(shape)
+        shape, static_shape = at.infer_static_shape(shape)
 
-        out = TensorType(dtype=a.type.dtype, shape=bcast)()
+        out = TensorType(dtype=a.type.dtype, shape=static_shape)()
 
         # Attempt to prevent in-place operations on this view-based output
         out.tag.indestructible = True
@@ -1637,11 +1670,14 @@ def grad(self, inputs, outputs_gradients):
         d_wrt_a = broadcast_to(dout, shape).sum(axis=new_dims)
 
         # Determine the dimensions that were broadcast
-        _, shape_bcast = at.infer_broadcastable(shape)
+        _, static_shape = at.infer_static_shape(shape)
+
+        # TODO: This needs to be performed at run-time when static shape
+        # information isn't available.
         bcast_sums = [
             i
-            for i, (a_b, s_b) in enumerate(zip(a.broadcastable, shape_bcast[-a.ndim :]))
-            if a_b and not s_b
+            for i, (a_s, s_s) in enumerate(zip(a.type.shape, static_shape[-a.ndim :]))
+            if a_s == 1 and s_s != 1
         ]
 
         if bcast_sums:

aesara/tensor/random/op.py

@@ -14,7 +14,7 @@
     constant,
     get_scalar_constant_value,
     get_vector_length,
-    infer_broadcastable,
+    infer_static_shape,
 )
 from aesara.tensor.random.type import RandomGeneratorType, RandomStateType, RandomType
 from aesara.tensor.random.utils import normalize_size_param, params_broadcast_shapes
@@ -322,7 +322,7 @@ def make_node(self, rng, size, dtype, *dist_params):
             )
 
         shape = self._infer_shape(size, dist_params)
-        _, bcast = infer_broadcastable(shape)
+        _, static_shape = infer_static_shape(shape)
         dtype = self.dtype or dtype
 
         if dtype == "floatX":
@@ -336,7 +336,7 @@ def make_node(self, rng, size, dtype, *dist_params):
             dtype_idx = constant(dtype, dtype="int64")
             dtype = all_dtypes[dtype_idx.data]
 
-        outtype = TensorType(dtype=dtype, shape=bcast)
+        outtype = TensorType(dtype=dtype, shape=static_shape)
         out_var = outtype()
         inputs = (rng, size, dtype_idx) + dist_params
         outputs = (rng.type(), out_var)

aesara/tensor/type.py

@@ -331,13 +331,6 @@ def convert_variable(self, var):
             # `specify_shape` will combine the more precise shapes of the two types
             return aesara.tensor.specify_shape(var, self.shape)
 
-    def value_zeros(self, shape):
-        """Create an numpy ndarray full of 0 values.
-
-        TODO: Remove this trivial method.
-        """
-        return np.zeros(shape, dtype=self.dtype)
-
     @staticmethod
     def values_eq(a, b, force_same_dtype=True):
         # TODO: check to see if the shapes must match; for now, we err on safe