Remove BroadcastTo

Author: ricardoV94

pytensor/link/jax/dispatch/extra_ops.py

@@ -3,10 +3,8 @@
 import jax.numpy as jnp
 
 from pytensor.link.jax.dispatch.basic import jax_funcify
-from pytensor.tensor.basic import infer_static_shape
 from pytensor.tensor.extra_ops import (
     Bartlett,
-    BroadcastTo,
     CumOp,
     FillDiagonal,
     FillDiagonalOffset,
@@ -102,18 +100,6 @@ def ravelmultiindex(*inp, mode=mode, order=order):
     return ravelmultiindex
 
 
-@jax_funcify.register(BroadcastTo)
-def jax_funcify_BroadcastTo(op, node, **kwargs):
-    shape = node.inputs[1:]
-    static_shape = infer_static_shape(shape)[1]
-
-    def broadcast_to(x, *shape):
-        shape = tuple(st if st is not None else s for s, st in zip(shape, static_shape))
-        return jnp.broadcast_to(x, shape)
-
-    return broadcast_to
-
-
 @jax_funcify.register(FillDiagonal)
 def jax_funcify_FillDiagonal(op, **kwargs):
     def filldiagonal(value, diagonal):

pytensor/link/numba/dispatch/extra_ops.py

@@ -2,15 +2,13 @@
 
 import numba
 import numpy as np
-from numba.misc.special import literal_unroll
 
 from pytensor import config
 from pytensor.link.numba.dispatch import basic as numba_basic
 from pytensor.link.numba.dispatch.basic import get_numba_type, numba_funcify
 from pytensor.raise_op import CheckAndRaise
 from pytensor.tensor.extra_ops import (
     Bartlett,
-    BroadcastTo,
     CumOp,
     FillDiagonal,
     FillDiagonalOffset,
@@ -353,29 +351,6 @@ def searchsorted(a, v):
     return searchsorted
 
 
-@numba_funcify.register(BroadcastTo)
-def numba_funcify_BroadcastTo(op, node, **kwargs):
-    create_zeros_tuple = numba_basic.create_tuple_creator(
-        lambda _: 0, len(node.inputs) - 1
-    )
-
-    # TODO broadcastable checks
-    @numba_basic.numba_njit
-    def broadcast_to(x, *shape):
-        scalars_shape = create_zeros_tuple()
-
-        i = 0
-        for s_i in literal_unroll(shape):
-            scalars_shape = numba_basic.tuple_setitem(
-                scalars_shape, i, numba_basic.to_scalar(s_i)
-            )
-            i += 1
-
-        return np.broadcast_to(x, scalars_shape)
-
-    return broadcast_to
-
-
 @numba_funcify.register(CheckAndRaise)
 def numba_funcify_CheckAndRaise(op, node, **kwargs):
     error = op.exc_type

pytensor/tensor/extra_ops.py

@@ -22,7 +22,7 @@
 from pytensor.scalar import int32 as int_t
 from pytensor.scalar import upcast
 from pytensor.tensor import basic as at
-from pytensor.tensor.basic import get_vector_length, second
+from pytensor.tensor.basic import alloc, second
 from pytensor.tensor.exceptions import NotScalarConstantError
 from pytensor.tensor.math import abs as at_abs
 from pytensor.tensor.math import all as pt_all
@@ -1561,141 +1561,6 @@ def broadcast_shape_iter(
     return tuple(result_dims)
 
 
-class BroadcastTo(COp):
-    """An `Op` for `numpy.broadcast_to`."""
-
-    _output_type_depends_on_input_value = True
-
-    __props__ = ()
-
-    view_map = {0: [0]}
-
-    def __call__(self, a, shape, **kwargs):
-        return super().__call__(a, *shape, **kwargs)
-
-    def make_node(self, a, *shape):
-        a = at.as_tensor_variable(a)
-
-        shape, static_shape = at.infer_static_shape(shape)
-
-        if len(shape) < a.ndim:
-            raise ValueError(
-                f"Broadcast target shape has {len(shape)} dims, which is shorter than input with {a.ndim} dims"
-            )
-
-        out = TensorType(dtype=a.type.dtype, shape=static_shape)()
-
-        # Attempt to prevent in-place operations on this view-based output
-        out.tag.indestructible = True
-
-        return Apply(self, [a] + shape, [out])
-
-    def perform(self, node, inputs, output_storage):
-        a, *shape = inputs
-        z = output_storage[0]
-        z[0] = np.broadcast_to(a, shape)
-
-    def grad(self, inputs, outputs_gradients):
-        a, *shape = inputs
-        (dout,) = outputs_gradients
-
-        # Determine the dimensions that were added by broadcasting
-        new_dims = list(range(dout.ndim - a.ndim))
-
-        d_wrt_a = broadcast_to(dout, shape).sum(axis=new_dims)
-
-        # Determine the dimensions that were broadcast
-        _, 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_s, s_s) in enumerate(zip(a.type.shape, static_shape[-a.ndim :]))
-            if a_s == 1 and s_s != 1
-        ]
-
-        if bcast_sums:
-            d_wrt_a = d_wrt_a.sum(axis=bcast_sums, keepdims=True)
-
-        return [d_wrt_a] + [
-            grad_undefined(self, i, shp) for i, shp in enumerate(shape, 1)
-        ]
-
-    def infer_shape(self, fgraph, node, ins_shapes):
-        return [node.inputs[1:]]
-
-    def c_code(self, node, name, inputs, outputs, sub):
-        inp_dims = node.inputs[0].ndim
-        out_dims = node.outputs[0].ndim
-        new_dims = out_dims - inp_dims
-
-        (x, *shape) = inputs
-        (out,) = outputs
-        fail = sub["fail"]
-
-        # TODO: Could just use `PyArray_Return`, no?
-        dims_array = ", ".join(
-            [
-                f"((dtype_{shape}*)(PyArray_DATA({shape})))[0]"
-                for i, shape in enumerate(shape)
-            ]
-        )
-
-        src = (
-            """
-            npy_intp itershape[%(out_dims)s] = {%(dims_array)s};
-
-            NpyIter *iter;
-            PyArrayObject *ops[1] = {%(x)s};
-            npy_uint32 flags = NPY_ITER_MULTI_INDEX | NPY_ITER_REFS_OK | NPY_ITER_ZEROSIZE_OK;
-            npy_uint32 op_flags[1] = {NPY_ITER_READONLY};
-            PyArray_Descr *op_dtypes[1] = {NULL};
-            int oa_ndim = %(out_dims)s;
-            int* op_axes[1] = {NULL};
-            npy_intp buffersize = 0;
-
-            for(int i = 0; i < %(inp_dims)s; i++)
-            {
-                if ((PyArray_DIMS(%(x)s)[i] != 1) && (PyArray_DIMS(%(x)s)[i] != itershape[i + %(new_dims)s]))
-                {
-                    PyErr_Format(PyExc_ValueError,
-                                 "Shape mismatch in broadcast_to: target shape[%%i] = %%lld is incompatible with input shape = %%lld.",
-                                 i,
-                                 (long long int) itershape[i + %(new_dims)s],
-                                 (long long int) PyArray_DIMS(%(x)s)[i]
-                    );
-                    %(fail)s
-                }
-            }
-
-            iter = NpyIter_AdvancedNew(
-                1, ops, flags, NPY_CORDER, NPY_NO_CASTING, op_flags, op_dtypes, oa_ndim, op_axes, itershape, buffersize
-            );
-            %(out)s = NpyIter_GetIterView(iter, 0);
-
-            if(%(out)s == NULL){
-                NpyIter_Deallocate(iter);
-                %(fail)s;
-            }
-
-            if (NpyIter_Deallocate(iter) != NPY_SUCCEED) {
-                %(fail)s;
-            }
-
-            """
-            % locals()
-        )
-
-        return src
-
-    def c_code_cache_version(self):
-        return (2,)
-
-
-broadcast_to_ = BroadcastTo()
-
-
 def geomspace(start, end, steps, base=10.0):
     from pytensor.tensor.math import log
 
@@ -1739,13 +1604,7 @@ def broadcast_to(
         broadcasted array may refer to a single memory location.
 
     """
-    x = at.as_tensor(x)
-    shape_len = get_vector_length(shape)
-
-    if x.ndim == 0 and shape_len == 0:
-        return x
-
-    return broadcast_to_(x, shape)
+    return alloc(x, *shape)
 
 
 def broadcast_arrays(*args: TensorVariable) -> Tuple[TensorVariable, ...]:

pytensor/tensor/rewriting/extra_ops.py

@@ -2,7 +2,7 @@
 from pytensor.graph.rewriting.basic import node_rewriter
 from pytensor.tensor.basic import Alloc, as_tensor_variable
 from pytensor.tensor.elemwise import Elemwise
-from pytensor.tensor.extra_ops import BroadcastTo, Repeat, Unique
+from pytensor.tensor.extra_ops import Repeat, Unique
 from pytensor.tensor.rewriting.basic import register_canonicalize, register_useless
 
 
@@ -60,39 +60,6 @@ def local_Unique_Alloc_lift(fgraph, node):
     return [new_x]
 
 
-@register_useless
-@register_canonicalize
-@node_rewriter([Unique])
-def local_Unique_BroadcastTo_lift(fgraph, node):
-    """Convert ``unique(broadcast_to(x, ...), axis=None)`` to ``unique(x, axis=None)``.
-
-    This isn't really so much a lift as a "reduction/consumption".
-    """
-    if not isinstance(node.op, Unique):
-        return False
-
-    if (
-        node.op.return_index
-        or node.op.return_inverse
-        or node.op.return_counts
-        or node.op.axis is not None
-    ):
-        return False
-
-    bcast_var = node.inputs[0]
-
-    if not (bcast_var.owner and isinstance(bcast_var.owner.op, BroadcastTo)):
-        return False
-
-    bcasted_var, *bcast_shape = bcast_var.owner.inputs
-
-    new_unique, *_ = node.op.make_node(bcasted_var).outputs
-
-    old_out = node.outputs[0]
-    new_x = as_tensor_variable(new_unique, ndim=old_out.ndim, dtype=old_out.dtype)
-    return [new_x]
-
-
 @register_useless
 @register_canonicalize
 @node_rewriter([Unique])
@@ -161,16 +128,3 @@ def local_Unique_second(fgraph, node):
     old_out = node.outputs[0]
     new_x = as_tensor_variable(new_unique, ndim=old_out.ndim, dtype=old_out.dtype)
     return [new_x]
-
-
-@register_useless
-@register_canonicalize
-@node_rewriter([BroadcastTo])
-def local_remove_scalar_BroadcastTo(fgraph, node):
-    bcast_shape = node.inputs[1:]
-
-    if not bcast_shape:
-        bcasted_var = node.inputs[0]
-        # If this isn't true, the graph is invalid
-        assert bcasted_var.ndim == 0
-        return [bcasted_var]

tests/link/jax/test_extra_ops.py

@@ -7,7 +7,7 @@
 from pytensor.graph.fg import FunctionGraph
 from pytensor.graph.op import get_test_value
 from pytensor.tensor import extra_ops as at_extra_ops
-from pytensor.tensor.type import matrix, vector
+from pytensor.tensor.type import matrix
 from tests.link.jax.test_basic import compare_jax_and_py
 
 
@@ -63,29 +63,6 @@ def test_extra_ops():
     )
 
 
-@pytest.mark.parametrize(
-    "x, shape",
-    [
-        (
-            set_test_value(
-                vector("x"), np.random.random(size=(2,)).astype(config.floatX)
-            ),
-            [at.as_tensor(3, dtype=np.int64), at.as_tensor(2, dtype=np.int64)],
-        ),
-        (
-            set_test_value(
-                vector("x"), np.random.random(size=(2,)).astype(config.floatX)
-            ),
-            [at.as_tensor(3, dtype=np.int8), at.as_tensor(2, dtype=np.int64)],
-        ),
-    ],
-)
-def test_BroadcastTo(x, shape):
-    out = at_extra_ops.broadcast_to(x, shape)
-    fgraph = FunctionGraph(outputs=[out])
-    compare_jax_and_py(fgraph, [get_test_value(i) for i in fgraph.inputs])
-
-
 @pytest.mark.xfail(
     version_parse(jax.__version__) >= version_parse("0.2.12"),
     reason="Omnistaging cannot be disabled",

tests/link/numba/test_extra_ops.py

@@ -36,41 +36,6 @@ def test_Bartlett(val):
     )
 
 
-@pytest.mark.parametrize(
-    "x, shape",
-    [
-        (
-            set_test_value(at.vector(), rng.random(size=(2,)).astype(config.floatX)),
-            [set_test_value(at.lscalar(), np.array(v)) for v in [3, 2]],
-        ),
-        (
-            set_test_value(at.vector(), rng.random(size=(2,)).astype(config.floatX)),
-            [at.as_tensor(3, dtype=np.int64), at.as_tensor(2, dtype=np.int64)],
-        ),
-        (
-            set_test_value(at.vector(), rng.random(size=(2,)).astype(config.floatX)),
-            at.as_tensor([set_test_value(at.lscalar(), np.array(v)) for v in [3, 2]]),
-        ),
-        (
-            set_test_value(at.vector(), rng.random(size=(2,)).astype(config.floatX)),
-            [at.as_tensor(3, dtype=np.int8), at.as_tensor(2, dtype=np.int64)],
-        ),
-    ],
-)
-def test_BroadcastTo(x, shape):
-    g = extra_ops.BroadcastTo()(x, shape)
-    g_fg = FunctionGraph(outputs=[g])
-
-    compare_numba_and_py(
-        g_fg,
-        [
-            i.tag.test_value
-            for i in g_fg.inputs
-            if not isinstance(i, (SharedVariable, Constant))
-        ],
-    )
-
-
 @pytest.mark.parametrize(
     "val, axis, mode",
     [