Remove Unbroadcast Op

Author: ricardoV94

Diff for: doc/library/tensor/basic.rst

@@ -619,9 +619,8 @@ dimensions, see :meth:`_tensor_py_operators.dimshuffle`.
 
 .. function:: shape_padleft(x, n_ones=1)
 
-    Reshape `x` by left padding the shape with `n_ones` 1s. Note that all
-    this new dimension will be broadcastable. To make them non-broadcastable
-    see the :func:`unbroadcast`.
+    Reshape `x` by left padding the shape with `n_ones` 1s.
+    All new dimensions will be broadcastable.
 
     :param x: variable to be reshaped
     :type x: any `TensorVariable` (or compatible)
@@ -633,9 +632,8 @@ dimensions, see :meth:`_tensor_py_operators.dimshuffle`.
 
 .. function:: shape_padright(x, n_ones=1)
 
-    Reshape `x` by right padding the shape with `n_ones` ones. Note that all
-    this new dimension will be broadcastable. To make them non-broadcastable
-    see the :func:`unbroadcast`.
+    Reshape `x` by right padding the shape with `n_ones` ones.
+    All new dimensions will be broadcastable.
 
     :param x: variable to be reshaped
     :type x: any TensorVariable (or compatible)
@@ -646,9 +644,8 @@ dimensions, see :meth:`_tensor_py_operators.dimshuffle`.
 
 .. function:: shape_padaxis(t, axis)
 
-    Reshape `t` by inserting ``1`` at the dimension `axis`. Note that this new
-    dimension will be broadcastable. To make it non-broadcastable
-    see the :func:`unbroadcast`.
+    Reshape `t` by inserting ``1`` at the dimension `axis`.
+    All new dimensions will be broadcastable.
 
     :type x: any `TensorVariable` (or compatible)
     :param x: variable to be reshaped

Diff for: pytensor/compile/function/pfunc.py

@@ -292,14 +292,8 @@ def clone_inputs(i):
                 f" shared_var.type={store_into.type},"
                 f" update_val={update_val}, update_val.type={getattr(update_val, 'type', None)})."
             )
-            err_sug = (
-                "If the difference is related to the broadcast pattern,"
-                " you can call the"
-                " tensor.shape.unbroadcast(var, axis_to_unbroadcast[, ...])"
-                " function to mask broadcastable dimensions."
-            )
 
-            raise TypeError(err_msg, err_sug)
+            raise TypeError(err_msg)
         assert store_into.type.is_super(update_val.type)
 
         update_d[store_into] = update_val

Diff for: pytensor/ifelse.py

@@ -26,7 +26,7 @@
 from pytensor.graph.replace import clone_replace
 from pytensor.graph.rewriting.basic import GraphRewriter, in2out, node_rewriter
 from pytensor.graph.type import HasDataType, HasShape
-from pytensor.tensor.shape import Reshape, Shape, SpecifyShape, Unbroadcast
+from pytensor.tensor.shape import Reshape, Shape, SpecifyShape
 
 
 if TYPE_CHECKING:
@@ -481,7 +481,6 @@ def cond_make_inplace(fgraph, node):
     Shape,
     SpecifyShape,
     Reshape,
-    Unbroadcast,
     pt.math.Dot,
     pt.math.Max,
     pt.math.Argmax,

Diff for: pytensor/link/jax/dispatch/shape.py

@@ -4,7 +4,7 @@
 from pytensor.graph.basic import Apply
 from pytensor.graph.op import Op
 from pytensor.link.jax.dispatch.basic import jax_funcify
-from pytensor.tensor.shape import Reshape, Shape, Shape_i, SpecifyShape, Unbroadcast
+from pytensor.tensor.shape import Reshape, Shape, Shape_i, SpecifyShape
 from pytensor.tensor.type import TensorType
 
 
@@ -104,11 +104,3 @@ def specifyshape(x, *shape):
         return x
 
     return specifyshape
-
-
-@jax_funcify.register(Unbroadcast)
-def jax_funcify_Unbroadcast(op, **kwargs):
-    def unbroadcast(x):
-        return x
-
-    return unbroadcast

Diff for: pytensor/link/numba/dispatch/tensor_basic.py

@@ -17,7 +17,6 @@
     Split,
     TensorFromScalar,
 )
-from pytensor.tensor.shape import Unbroadcast
 
 
 @numba_funcify.register(AllocEmpty)
@@ -232,15 +231,6 @@ def makevector({", ".join(input_names)}):
     return numba_basic.numba_njit(makevector_fn)
 
 
-@numba_funcify.register(Unbroadcast)
-def numba_funcify_Unbroadcast(op, **kwargs):
-    @numba_basic.numba_njit
-    def unbroadcast(x):
-        return x
-
-    return unbroadcast
-
-
 @numba_funcify.register(TensorFromScalar)
 def numba_funcify_TensorFromScalar(op, **kwargs):
     @numba_basic.numba_njit(inline="always")

Diff for: pytensor/link/pytorch/dispatch/shape.py

@@ -2,7 +2,7 @@
 
 from pytensor.graph.basic import Constant
 from pytensor.link.pytorch.dispatch.basic import pytorch_funcify
-from pytensor.tensor.shape import Reshape, Shape, Shape_i, SpecifyShape, Unbroadcast
+from pytensor.tensor.shape import Reshape, Shape, Shape_i, SpecifyShape
 
 
 @pytorch_funcify.register(Reshape)
@@ -56,11 +56,3 @@ def specifyshape(x, *shape):
         return x
 
     return specifyshape
-
-
-@pytorch_funcify.register(Unbroadcast)
-def pytorch_funcify_Unbroadcast(op, **kwargs):
-    def unbroadcast(x):
-        return x
-
-    return unbroadcast

Diff for: pytensor/scan/basic.py

@@ -15,7 +15,7 @@
 from pytensor.tensor.basic import get_underlying_scalar_constant_value
 from pytensor.tensor.exceptions import NotScalarConstantError
 from pytensor.tensor.math import minimum
-from pytensor.tensor.shape import shape_padleft, unbroadcast
+from pytensor.tensor.shape import shape_padleft
 from pytensor.tensor.type import TensorType, integer_dtypes
 from pytensor.updates import OrderedUpdates
 
@@ -748,7 +748,7 @@ def wrap_into_list(x):
             # defined in scan utils
             sit_sot_scan_inputs.append(
                 expand_empty(
-                    unbroadcast(shape_padleft(actual_arg), 0),
+                    shape_padleft(actual_arg),
                     actual_n_steps,
                 )
             )
@@ -865,13 +865,13 @@ def wrap_into_list(x):
     if n_fixed_steps in (1, -1):
         for pos, inner_out in enumerate(outputs):
             # we need to see if we need to pad our sequences with an
-            # unbroadcastable dimension; case example : we return an
+            # extra dimension; case example : we return an
             # output for which we want all intermediate. If n_steps is 1
             # then, if we return the output as given by the innner function
             # this will represent only a slice and it will have one
             # dimension less.
             if isinstance(inner_out.type, TensorType) and return_steps.get(pos, 0) != 1:
-                outputs[pos] = unbroadcast(shape_padleft(inner_out), 0)
+                outputs[pos] = shape_padleft(inner_out)
 
         if not return_list and len(outputs) == 1:
             outputs = outputs[0]
@@ -1002,7 +1002,7 @@ def wrap_into_list(x):
                 sit_sot_inner_inputs.append(new_var)
                 sit_sot_scan_inputs.append(
                     expand_empty(
-                        unbroadcast(shape_padleft(input.variable), 0),
+                        shape_padleft(input.variable),
                         actual_n_steps,
                     )
                 )

Diff for: pytensor/scan/op.py

@@ -166,8 +166,7 @@ def check_broadcast(v1, v2):
         "axis %d in `output_info`. This can happen if one of the "
         "dimension is fixed to 1 in the input, while it is still "
         "variable in the output, or vice-verca. You have to make "
-        "them consistent, e.g. using pytensor.tensor."
-        "{unbroadcast, specify_broadcastable}."
+        "them consistent, e.g. using pytensor.tensor.specify_broadcastable."
     )
     size = min(v1.type.ndim, v2.type.ndim)
     for n, (b1, b2) in enumerate(

Diff for: pytensor/tensor/basic.py

@@ -53,7 +53,6 @@
 from pytensor.tensor.shape import (
     Shape,
     Shape_i,
-    Unbroadcast,
     shape,
     shape_padaxis,
     shape_padleft,
@@ -334,9 +333,7 @@ def _get_underlying_scalar_constant_value(
         if not only_process_constants and getattr(v, "owner", None) and max_recur > 0:
             op = v.owner.op
             max_recur -= 1
-            if isinstance(
-                op, Alloc | DimShuffle | Unbroadcast | OutputGuard | DeepCopyOp
-            ):
+            if isinstance(op, Alloc | DimShuffle | OutputGuard | DeepCopyOp):
                 # OutputGuard is only used in debugmode but we
                 # keep it here to avoid problems with old pickles
                 v = v.owner.inputs[0]
@@ -498,14 +495,6 @@ def _get_underlying_scalar_constant_value(
                     grandparent = leftmost_parent.owner.inputs[0]
                     gp_shape = grandparent.type.shape
                     ndim = grandparent.type.ndim
-                    if grandparent.owner and isinstance(
-                        grandparent.owner.op, Unbroadcast
-                    ):
-                        ggp_shape = grandparent.owner.inputs[0].type.shape
-                        l = [
-                            _get_underlying_scalar_constant_value(s) for s in ggp_shape
-                        ]
-                        gp_shape = tuple(l)
 
                     if not (idx < ndim):
                         msg = (

Diff for: pytensor/tensor/rewriting/shape.py

@@ -42,9 +42,7 @@
     Shape,
     Shape_i,
     SpecifyShape,
-    Unbroadcast,
     specify_shape,
-    unbroadcast,
 )
 from pytensor.tensor.subtensor import Subtensor, get_idx_list
 from pytensor.tensor.type import TensorType, discrete_dtypes, integer_dtypes
@@ -1296,78 +1294,3 @@ def local_track_shape_i(fgraph, node):
     # structure.
     replacement = shape_feature.scheduled[node]
     return [shape_feature.shape_of[replacement][node.op.i]]
-
-
-@register_useless
-@register_canonicalize
-@register_specialize
-@node_rewriter([Unbroadcast])
-def local_useless_unbroadcast(fgraph, node):
-    """Remove `Unbroadcast` if it does not actually change the broadcasting pattern."""
-    if isinstance(node.op, Unbroadcast):
-        x = node.inputs[0]
-        if x.type.ndim == node.outputs[0].type.ndim and all(
-            s1 == s2
-            for s1, s2 in zip(x.type.shape, node.outputs[0].type.shape, strict=True)
-            if s1 == 1 or s2 == 1
-        ):
-            # No broadcastable flag was modified
-            # No need to copy over stack trace,
-            # because x should already have a stack trace.
-            return [x]
-        else:
-            # Keep the flags that modify something
-            new_axes = tuple(ax for ax in node.op.axes if x.type.shape[ax] == 1)
-            if new_axes == node.op.axes:
-                # All flags are useful
-                return None
-            else:
-                r = unbroadcast(x, *new_axes)
-                # Copy over stacktrace from previous output
-                copy_stack_trace(node.outputs, r)
-                return [r]
-
-
-@register_canonicalize
-@register_specialize
-@node_rewriter([Unbroadcast])
-def local_unbroadcast_lift(fgraph, node):
-    """
-    Lifts `Unbroadcast` through unary Elemwise operations,
-    and merges consecutive `Unbroadcast`s.
-
-    Unbroadcast(Elemwise(x)) => Elemwise(Unbroadcast(x))
-    Unbroadcast(Unbroadcast(x)) => Unbroadcast(x)
-
-    TODO: Implement equivalent Elemwise lift for SpecifyShape
-    """
-    op = node.op
-    if not isinstance(op, Unbroadcast):
-        return False
-
-    inp = node.inputs[0]
-    inode = inp.owner
-    if inode and isinstance(inode.op, Elemwise) and len(inode.inputs) == 1:
-        if len(fgraph.clients.get(inp, ())) == 1:
-            unbroadcasted = unbroadcast(inode.inputs[0], *op.axes)
-            copy_stack_trace(node.outputs, unbroadcasted)
-
-            rval = inode.op.make_node(unbroadcasted).outputs
-
-            # Copy over stacktrace from previous output (after unbroadcasting)
-            # and input (after elemwise operation) to new output, because an
-            # error in the new graph could have been caused by either of the
-            # two ops.
-            copy_stack_trace(node.outputs + node.inputs, rval)
-            return rval
-
-    if inode and isinstance(inode.op, Unbroadcast):
-        # Merge axis of each unbroadcast
-        axis = tuple(set(inode.op.axes).union(set(op.axes)))
-        iinput = inode.inputs[0]
-        rval = [unbroadcast(iinput, *axis)]
-        # Copy over stacktrace from previous output (after second unbroadcasting)
-        # and from previous input (after first unbroadcasting) because an error in
-        # the new graph could have been caused by either of the two Unbroadcast ops.
-        copy_stack_trace(node.outputs + node.inputs, rval)
-        return rval

Diff for: pytensor/tensor/rewriting/subtensor.py

@@ -59,11 +59,9 @@
 from pytensor.tensor.shape import (
     Shape,
     SpecifyShape,
-    Unbroadcast,
     shape_padleft,
     shape_tuple,
     specify_shape,
-    unbroadcast,
 )
 from pytensor.tensor.sharedvar import TensorSharedVariable
 from pytensor.tensor.subtensor import (
@@ -429,7 +427,6 @@ def local_subtensor_lift(fgraph, node):
     Handles the following unary ops:
     elemwise(x,...)[idx] -> elemwise(x[idx],...)
       when x,... are broadcasted scalar or not broadcasted at all
-    Unbroadcast(x)[idx] => Unbroadcast(x[idx])
 
     """
     if isinstance(node.op, Subtensor):
@@ -488,40 +485,6 @@ def local_subtensor_lift(fgraph, node):
                 copy_stack_trace([node.outputs[0], node.inputs[0]], ret)
                 return [ret]
 
-        if isinstance(u.owner.op, Unbroadcast):
-            # Subtensor might reduce dim., adapt broadcast pattern accordingly
-            old_axes = u.owner.op.axes
-            new_axes = []
-
-            # loop through indices being subtensor-ed
-            # i indexes broadcastable pattern before subtensor
-            # j indexes broadcastable pattern after subtensor
-            j = 0
-            for i, x in enumerate(node.op.idx_list):
-                # if it is not a slice, it will reduce the dimension, should
-                # not appear in the broascastable dimensions
-                if isinstance(x, slice):
-                    if i in old_axes:
-                        new_axes.append(j)
-                    j += 1
-            # now keep the broadcastable pattern of all
-            # items not appearing in subtensor list
-            for i in range(len(node.op.idx_list), len(u.broadcastable)):
-                if i in old_axes:
-                    new_axes.append(j)
-                j += 1
-
-            subt_x = node.op(u.owner.inputs[0], *node.inputs[1:])
-            # Copy over previous output stacktrace
-            copy_stack_trace(node.outputs[0], subt_x)
-
-            rbcast_subt_x = unbroadcast(subt_x, *new_axes)
-            # Copy over previous output stacktrace
-            # and stacktrace from previous unary operation
-            copy_stack_trace([node.outputs[0], node.inputs[0]], rbcast_subt_x)
-
-            return [rbcast_subt_x]
-
 
 @register_canonicalize
 @register_specialize