pytensor 2.28.2__tar.gz → 2.29.0__tar.gz

{pytensor-2.28.2/pytensor.egg-info → pytensor-2.29.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: pytensor
-Version: 2.28.2
+Version: 2.29.0
 Summary: Optimizing compiler for evaluating mathematical expressions on CPUs and GPUs.
 Author-email: pymc-devs <pymc.devs@gmail.com>
 License: .. _license:

{pytensor-2.28.2 → pytensor-2.29.0}/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

{pytensor-2.28.2 → pytensor-2.29.0}/pytensor/_version.py

@@ -8,11 +8,11 @@ import json
 
 version_json = '''
 {
- "date": "2025-03-03T18:16:56+0100",
+ "date": "2025-03-18T14:02:26+0100",
  "dirty": false,
  "error": null,
- "full-revisionid": "757a10cd496ac24800f03491891ff6bba2d4807a",
- "version": "2.28.2"
+ "full-revisionid": "a149f6c996bd0499406150853bef517e7d7a3c4f",
+ "version": "2.29.0"
 }
 '''  # END VERSION_JSON
 

{pytensor-2.28.2 → pytensor-2.29.0}/pytensor/compile/function/pfunc.py

@@ -292,14 +292,8 @@ def rebuild_collect_shared(
                 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

{pytensor-2.28.2 → pytensor-2.29.0}/pytensor/compile/mode.py

@@ -454,6 +454,19 @@ else:
         RewriteDatabaseQuery(include=["fast_run", "py_only"]),
     )
 
+NUMBA = Mode(
+    NumbaLinker(),
+    RewriteDatabaseQuery(
+        include=["fast_run", "numba"],
+        exclude=[
+            "cxx_only",
+            "BlasOpt",
+            "local_careduce_fusion",
+            "scan_save_mem_prealloc",
+        ],
+    ),
+)
+
 JAX = Mode(
     JAXLinker(),
     RewriteDatabaseQuery(
@@ -463,6 +476,7 @@ JAX = Mode(
             "BlasOpt",
             "fusion",
             "inplace",
+            "scan_save_mem_prealloc",
         ],
     ),
 )
@@ -476,16 +490,10 @@ PYTORCH = Mode(
             "fusion",
             "inplace",
             "local_uint_constant_indices",
+            "scan_save_mem_prealloc",
         ],
     ),
 )
-NUMBA = Mode(
-    NumbaLinker(),
-    RewriteDatabaseQuery(
-        include=["fast_run", "numba"],
-        exclude=["cxx_only", "BlasOpt", "local_careduce_fusion"],
-    ),
-)
 
 
 predefined_modes = {

{pytensor-2.28.2 → pytensor-2.29.0}/pytensor/configdefaults.py

@@ -1085,7 +1085,9 @@ def add_scan_configvars():
         "scan__allow_output_prealloc",
         "Allow/disallow memory preallocation for outputs inside of scan "
         "(default: True)",
-        BoolParam(True),
+        # Non-mutable because ScanSaveMem rewrite checks it,
+        # and we can't have the rewrite and the implementation mismatch
+        BoolParam(True, mutable=False),
         in_c_key=False,
     )
 

{pytensor-2.28.2 → pytensor-2.29.0}/pytensor/ifelse.py

@@ -26,7 +26,7 @@ from pytensor.graph.op import _NoPythonOp
 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 @@ acceptable_ops = (
     Shape,
     SpecifyShape,
     Reshape,
-    Unbroadcast,
     pt.math.Dot,
     pt.math.Max,
     pt.math.Argmax,

{pytensor-2.28.2 → pytensor-2.29.0}/pytensor/link/jax/dispatch/scan.py

@@ -29,7 +29,7 @@ def jax_funcify_Scan(op: Scan, **kwargs):
         # Extract JAX scan inputs
         outer_inputs = list(outer_inputs)
         n_steps = outer_inputs[0]  # JAX `length`
-        seqs = op.outer_seqs(outer_inputs)  # JAX `xs`
+        seqs = [seq[:n_steps] for seq in op.outer_seqs(outer_inputs)]  # JAX `xs`
 
         mit_sot_init = []
         for tap, seq in zip(

{pytensor-2.28.2 → pytensor-2.29.0}/pytensor/link/jax/dispatch/shape.py

@@ -4,7 +4,7 @@ from pytensor.graph import Constant
 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 jax_funcify_SpecifyShape(op, node, **kwargs):
         return x
 
     return specifyshape
-
-
-@jax_funcify.register(Unbroadcast)
-def jax_funcify_Unbroadcast(op, **kwargs):
-    def unbroadcast(x):
-        return x
-
-    return unbroadcast

{pytensor-2.28.2 → pytensor-2.29.0}/pytensor/link/jax/dispatch/slinalg.py

@@ -1,3 +1,5 @@
+import warnings
+
 import jax
 
 from pytensor.link.jax.dispatch.basic import jax_funcify
@@ -39,13 +41,29 @@ def jax_funcify_Cholesky(op, **kwargs):
 
 @jax_funcify.register(Solve)
 def jax_funcify_Solve(op, **kwargs):
-    if op.assume_a != "gen" and op.lower:
-        lower = True
+    assume_a = op.assume_a
+    lower = op.lower
+
+    if assume_a == "tridiagonal":
+        # jax.scipy.solve does not yet support tridiagonal matrices
+        # But there's a jax.lax.linalg.tridiaonal_solve we can use instead.
+        def solve(a, b):
+            dl = jax.numpy.diagonal(a, offset=-1, axis1=-2, axis2=-1)
+            d = jax.numpy.diagonal(a, offset=0, axis1=-2, axis2=-1)
+            du = jax.numpy.diagonal(a, offset=1, axis1=-2, axis2=-1)
+            return jax.lax.linalg.tridiagonal_solve(dl, d, du, b, lower=lower)
+
     else:
-        lower = False
+        if assume_a not in ("gen", "sym", "her", "pos"):
+            warnings.warn(
+                f"JAX solve does not support assume_a={op.assume_a}. Defaulting to assume_a='gen'.\n"
+                f"If appropriate, you may want to set assume_a to one of 'sym', 'pos', 'her' or 'tridiagonal' to improve performance.",
+                UserWarning,
+            )
+            assume_a = "gen"
 
-    def solve(a, b, lower=lower):
-        return jax.scipy.linalg.solve(a, b, lower=lower)
+        def solve(a, b):
+            return jax.scipy.linalg.solve(a, b, lower=lower, assume_a=assume_a)
 
     return solve
 
@@ -53,7 +71,6 @@ def jax_funcify_Solve(op, **kwargs):
 @jax_funcify.register(SolveTriangular)
 def jax_funcify_SolveTriangular(op, **kwargs):
     lower = op.lower
-    trans = op.trans
     unit_diagonal = op.unit_diagonal
     check_finite = op.check_finite
 
@@ -62,7 +79,7 @@ def jax_funcify_SolveTriangular(op, **kwargs):
             A,
             b,
             lower=lower,
-            trans=trans,
+            trans=0,  # this is handled by explicitly transposing A, so it will always be 0 when we get to here.
             unit_diagonal=unit_diagonal,
             check_finite=check_finite,
         )

{pytensor-2.28.2 → pytensor-2.29.0}/pytensor/link/numba/dispatch/basic.py

@@ -243,93 +243,6 @@ def impl_to_scalar(x):
         raise TypingError(f"{x} must be a scalar compatible type.")
 
 
-def enable_slice_literals():
-    """Enable lowering for ``SliceLiteral``s.
-
-    TODO: This can be removed once https://github.com/numba/numba/pull/6996 is merged
-    and a release is made.
-    """
-    from numba.core import types
-    from numba.core.datamodel.models import SliceModel
-    from numba.core.datamodel.registry import register_default
-    from numba.core.imputils import lower_cast, lower_constant
-    from numba.core.types.misc import SliceLiteral
-    from numba.cpython.slicing import get_defaults
-
-    register_default(numba.types.misc.SliceLiteral)(SliceModel)
-
-    @property
-    def key(self):
-        return self.name
-
-    SliceLiteral.key = key
-
-    def make_slice_from_constant(context, builder, ty, pyval):
-        sli = context.make_helper(builder, ty)
-        lty = context.get_value_type(types.intp)
-
-        (
-            default_start_pos,
-            default_start_neg,
-            default_stop_pos,
-            default_stop_neg,
-            default_step,
-        ) = (context.get_constant(types.intp, x) for x in get_defaults(context))
-
-        step = pyval.step
-        if step is None:
-            step_is_neg = False
-            step = default_step
-        else:
-            step_is_neg = step < 0
-            step = lty(step)
-
-        start = pyval.start
-        if start is None:
-            if step_is_neg:
-                start = default_start_neg
-            else:
-                start = default_start_pos
-        else:
-            start = lty(start)
-
-        stop = pyval.stop
-        if stop is None:
-            if step_is_neg:
-                stop = default_stop_neg
-            else:
-                stop = default_stop_pos
-        else:
-            stop = lty(stop)
-
-        sli.start = start
-        sli.stop = stop
-        sli.step = step
-
-        return sli._getvalue()
-
-    @lower_constant(numba.types.SliceType)
-    def constant_slice(context, builder, ty, pyval):
-        if isinstance(ty, types.Literal):
-            typ = ty.literal_type
-        else:
-            typ = ty
-
-        return make_slice_from_constant(context, builder, typ, pyval)
-
-    @lower_cast(numba.types.misc.SliceLiteral, numba.types.SliceType)
-    def cast_from_literal(context, builder, fromty, toty, val):
-        return make_slice_from_constant(
-            context,
-            builder,
-            toty,
-            fromty.literal_value,
-        )
-
-
-enable_slice_literals()
-
-
 def create_tuple_creator(f, n):
     """Construct a compile-time ``tuple``-comprehension-like loop.
 

{pytensor-2.28.2 → pytensor-2.29.0}/pytensor/link/numba/dispatch/scan.py

@@ -55,7 +55,7 @@ def array0d_range(x):
 
 
 @numba_funcify.register(Scan)
-def numba_funcify_Scan(op, node, **kwargs):
+def numba_funcify_Scan(op: Scan, node, **kwargs):
     # Apply inner rewrites
     # TODO: Not sure this is the right place to do this, should we have a rewrite that
     #  explicitly triggers the optimization of the inner graphs of Scan?
@@ -67,9 +67,32 @@ def numba_funcify_Scan(op, node, **kwargs):
         .optimizer
     )
     fgraph = op.fgraph
+    # When the buffer can only hold one SITSOT or as as many MITSOT as there are taps,
+    # We must always discard the oldest tap, so it's safe to destroy it in the inner function.
+    # TODO: Allow inplace for MITMOT
+    destroyable_sitsot = [
+        inner_sitsot
+        for outer_sitsot, inner_sitsot in zip(
+            op.outer_sitsot(node.inputs), op.inner_sitsot(fgraph.inputs), strict=True
+        )
+        if outer_sitsot.type.shape[0] == 1
+    ]
+    destroyable_mitsot = [
+        oldest_inner_mitmot
+        for outer_mitsot, oldest_inner_mitmot, taps in zip(
+            op.outer_mitsot(node.inputs),
+            op.oldest_inner_mitsot(fgraph.inputs),
+            op.info.mit_sot_in_slices,
+            strict=True,
+        )
+        if outer_mitsot.type.shape[0] == abs(min(taps))
+    ]
+    destroyable = {*destroyable_sitsot, *destroyable_mitsot}
     add_supervisor_to_fgraph(
         fgraph=fgraph,
-        input_specs=[In(x, borrow=True, mutable=False) for x in fgraph.inputs],
+        input_specs=[
+            In(x, borrow=True, mutable=x in destroyable) for x in fgraph.inputs
+        ],
         accept_inplace=True,
     )
     rewriter(fgraph)
@@ -222,14 +245,16 @@ def numba_funcify_Scan(op, node, **kwargs):
         # the storage array.
         # This is needed when the output storage array does not have a length
         # equal to the number of taps plus `n_steps`.
+        # If the storage size only allows one entry, there's nothing to rotate
         output_storage_post_proc_stmts.append(
             dedent(
                 f"""
-                if (i + {tap_size}) > {storage_size}:
+                if 1 < {storage_size} < (i + {tap_size}):
                     {outer_in_name}_shift = (i + {tap_size}) % ({storage_size})
-                    {outer_in_name}_left = {outer_in_name}[:{outer_in_name}_shift]
-                    {outer_in_name}_right = {outer_in_name}[{outer_in_name}_shift:]
-                    {outer_in_name} = np.concatenate(({outer_in_name}_right, {outer_in_name}_left))
+                    if {outer_in_name}_shift > 0:
+                        {outer_in_name}_left = {outer_in_name}[:{outer_in_name}_shift]
+                        {outer_in_name}_right = {outer_in_name}[{outer_in_name}_shift:]
+                        {outer_in_name} = np.concatenate(({outer_in_name}_right, {outer_in_name}_left))
                 """
             ).strip()
         )
@@ -417,4 +442,4 @@ def scan({", ".join(outer_in_names)}):
 
     scan_op_fn = compile_function_src(scan_op_src, "scan", {**globals(), **global_env})
 
-    return numba_basic.numba_njit(scan_op_fn)
+    return numba_basic.numba_njit(scan_op_fn, boundscheck=False)

{pytensor-2.28.2 → pytensor-2.29.0}/pytensor/link/numba/dispatch/slinalg.py

@@ -1,3 +1,4 @@
+import warnings
 from collections.abc import Callable
 
 import numba
@@ -180,7 +181,6 @@ def solve_triangular_impl(A, B, trans, lower, unit_diagonal, b_ndim, overwrite_b
 
 @numba_funcify.register(SolveTriangular)
 def numba_funcify_SolveTriangular(op, node, **kwargs):
-    trans = bool(op.trans)
     lower = op.lower
     unit_diagonal = op.unit_diagonal
     check_finite = op.check_finite
@@ -208,7 +208,7 @@ def numba_funcify_SolveTriangular(op, node, **kwargs):
         res = _solve_triangular(
             a,
             b,
-            trans=trans,
+            trans=0,  # transposing is handled explicitly on the graph, so we never use this argument
             lower=lower,
             unit_diagonal=unit_diagonal,
             overwrite_b=overwrite_b,
@@ -654,7 +654,7 @@ def solve_gen_impl(
 
 def _sysv(
     A: np.ndarray, B: np.ndarray, lower: bool, overwrite_a: bool, overwrite_b: bool
-) -> tuple[np.ndarray, np.ndarray, int]:
+) -> tuple[np.ndarray, np.ndarray, np.ndarray, int]:
     """
     Placeholder for solving a linear system with a symmetric matrix; used by linalg.solve.
     """
@@ -665,7 +665,8 @@ def _sysv(
 def sysv_impl(
     A: np.ndarray, B: np.ndarray, lower: bool, overwrite_a: bool, overwrite_b: bool
 ) -> Callable[
-    [np.ndarray, np.ndarray, bool, bool, bool], tuple[np.ndarray, np.ndarray, int]
+    [np.ndarray, np.ndarray, bool, bool, bool],
+    tuple[np.ndarray, np.ndarray, np.ndarray, int],
 ]:
     ensure_lapack()
     _check_scipy_linalg_matrix(A, "sysv")
@@ -741,8 +742,8 @@ def sysv_impl(
         )
 
         if B_is_1d:
-            return B_copy[..., 0], IPIV, int_ptr_to_val(INFO)
-        return B_copy, IPIV, int_ptr_to_val(INFO)
+            B_copy = B_copy[..., 0]
+        return A_copy, B_copy, IPIV, int_ptr_to_val(INFO)
 
     return impl
 
@@ -771,7 +772,7 @@ def sycon_impl(
 
         N = val_to_int_ptr(_N)
         LDA = val_to_int_ptr(_N)
-        UPLO = val_to_int_ptr(ord("L"))
+        UPLO = val_to_int_ptr(ord("U"))
         ANORM = np.array(anorm, dtype=dtype)
         RCOND = np.empty(1, dtype=dtype)
         WORK = np.empty(2 * _N, dtype=dtype)
@@ -844,10 +845,10 @@ def solve_symmetric_impl(
     ) -> np.ndarray:
         _solve_check_input_shapes(A, B)
 
-        x, ipiv, info = _sysv(A, B, lower, overwrite_a, overwrite_b)
+        lu, x, ipiv, info = _sysv(A, B, lower, overwrite_a, overwrite_b)
         _solve_check(A.shape[-1], info)
 
-        rcond, info = _sycon(A, ipiv, _xlange(A, order="I"))
+        rcond, info = _sycon(lu, ipiv, _xlange(A, order="I"))
         _solve_check(A.shape[-1], info, True, rcond)
 
         return x
@@ -1071,14 +1072,17 @@ def numba_funcify_Solve(op, node, **kwargs):
     elif assume_a == "sym":
         solve_fn = _solve_symmetric
     elif assume_a == "her":
-        raise NotImplementedError(
-            'Use assume_a = "sym" for symmetric real matrices. If you need compelx support, '
-            "please open an issue on github."
-        )
+        # We already ruled out complex inputs
+        solve_fn = _solve_symmetric
     elif assume_a == "pos":
         solve_fn = _solve_psd
     else:
-        raise NotImplementedError(f"Assumption {assume_a} not supported in Numba mode")
+        warnings.warn(
+            f"Numba assume_a={assume_a} not implemented. Falling back to general solve.\n"
+            f"If appropriate, you may want to set assume_a to one of 'sym', 'pos', or 'her' to improve performance.",
+            UserWarning,
+        )
+        solve_fn = _solve_gen
 
     @numba_basic.numba_njit(inline="always")
     def solve(a, b):

{pytensor-2.28.2 → pytensor-2.29.0}/pytensor/link/numba/dispatch/subtensor.py

@@ -130,15 +130,6 @@ def numba_funcify_AdvancedSubtensor(op, node, **kwargs):
         if isinstance(idx.type, TensorType)
     ]
 
-    def broadcasted_to(x_bcast: tuple[bool, ...], to_bcast: tuple[bool, ...]):
-        # Check that x is not broadcasted to y based on broadcastable info
-        if len(x_bcast) < len(to_bcast):
-            return True
-        for x_bcast_dim, to_bcast_dim in zip(x_bcast, to_bcast, strict=True):
-            if x_bcast_dim and not to_bcast_dim:
-                return True
-        return False
-
     # Special implementation for consecutive integer vector indices
     if (
         not basic_idxs
@@ -151,17 +142,6 @@ def numba_funcify_AdvancedSubtensor(op, node, **kwargs):
         )
         # Must be consecutive
         and not op.non_consecutive_adv_indexing(node)
-        # y in set/inc_subtensor cannot be broadcasted
-        and (
-            y is None
-            or not broadcasted_to(
-                y.type.broadcastable,
-                (
-                    x.type.broadcastable[: adv_idxs[0]["axis"]]
-                    + x.type.broadcastable[adv_idxs[-1]["axis"] :]
-                ),
-            )
-        )
     ):
         return numba_funcify_multiple_integer_vector_indexing(op, node, **kwargs)
 
@@ -191,14 +171,24 @@ def numba_funcify_AdvancedSubtensor(op, node, **kwargs):
     return numba_funcify_default_subtensor(op, node, **kwargs)
 
 
+def _broadcasted_to(x_bcast: tuple[bool, ...], to_bcast: tuple[bool, ...]):
+    # Check that x is not broadcasted to y based on broadcastable info
+    if len(x_bcast) < len(to_bcast):
+        return True
+    for x_bcast_dim, to_bcast_dim in zip(x_bcast, to_bcast, strict=True):
+        if x_bcast_dim and not to_bcast_dim:
+            return True
+    return False
+
+
 def numba_funcify_multiple_integer_vector_indexing(
     op: AdvancedSubtensor | AdvancedIncSubtensor, node, **kwargs
 ):
     # Special-case implementation for multiple consecutive vector integer indices (and set/incsubtensor)
     if isinstance(op, AdvancedSubtensor):
-        y, idxs = None, node.inputs[1:]
+        idxs = node.inputs[1:]
     else:
-        y, *idxs = node.inputs[1:]
+        idxs = node.inputs[2:]
 
     first_axis = next(
         i for i, idx in enumerate(idxs) if isinstance(idx.type, TensorType)
@@ -211,6 +201,10 @@ def numba_funcify_multiple_integer_vector_indexing(
         )
     except StopIteration:
         after_last_axis = len(idxs)
+    last_axis = after_last_axis - 1
+
+    vector_indices = idxs[first_axis:after_last_axis]
+    assert all(v.type.broadcastable == (False,) for v in vector_indices)
 
     if isinstance(op, AdvancedSubtensor):
 
@@ -231,43 +225,59 @@ def numba_funcify_multiple_integer_vector_indexing(
 
         return advanced_subtensor_multiple_vector
 
-    elif op.set_instead_of_inc:
+    else:
         inplace = op.inplace
 
-        @numba_njit
-        def advanced_set_subtensor_multiple_vector(x, y, *idxs):
-            vec_idxs = idxs[first_axis:after_last_axis]
-            x_shape = x.shape
+        # Check if y must be broadcasted
+        # Includes the last integer vector index,
+        x, y = node.inputs[:2]
+        indexed_bcast_dims = (
+            *x.type.broadcastable[:first_axis],
+            *x.type.broadcastable[last_axis:],
+        )
+        y_is_broadcasted = _broadcasted_to(y.type.broadcastable, indexed_bcast_dims)
 
-            if inplace:
-                out = x
-            else:
-                out = x.copy()
+        if op.set_instead_of_inc:
 
-            for outer in np.ndindex(x_shape[:first_axis]):
-                for i, scalar_idxs in enumerate(zip(*vec_idxs)):  # noqa: B905
-                    out[(*outer, *scalar_idxs)] = y[(*outer, i)]
-            return out
+            @numba_njit
+            def advanced_set_subtensor_multiple_vector(x, y, *idxs):
+                vec_idxs = idxs[first_axis:after_last_axis]
+                x_shape = x.shape
 
-        return advanced_set_subtensor_multiple_vector
+                if inplace:
+                    out = x
+                else:
+                    out = x.copy()
 
-    else:
-        inplace = op.inplace
+                if y_is_broadcasted:
+                    y = np.broadcast_to(y, x_shape[:first_axis] + x_shape[last_axis:])
 
-        @numba_njit
-        def advanced_inc_subtensor_multiple_vector(x, y, *idxs):
-            vec_idxs = idxs[first_axis:after_last_axis]
-            x_shape = x.shape
+                for outer in np.ndindex(x_shape[:first_axis]):
+                    for i, scalar_idxs in enumerate(zip(*vec_idxs)):  # noqa: B905
+                        out[(*outer, *scalar_idxs)] = y[(*outer, i)]
+                return out
+
+            return advanced_set_subtensor_multiple_vector
+
+        else:
+
+            @numba_njit
+            def advanced_inc_subtensor_multiple_vector(x, y, *idxs):
+                vec_idxs = idxs[first_axis:after_last_axis]
+                x_shape = x.shape
+
+                if inplace:
+                    out = x
+                else:
+                    out = x.copy()
 
-            if inplace:
-                out = x
-            else:
-                out = x.copy()
+                if y_is_broadcasted:
+                    y = np.broadcast_to(y, x_shape[:first_axis] + x_shape[last_axis:])
 
-            for outer in np.ndindex(x_shape[:first_axis]):
-                for i, scalar_idxs in enumerate(zip(*vec_idxs)):  # noqa: B905
-                    out[(*outer, *scalar_idxs)] += y[(*outer, i)]
-            return out
+                for outer in np.ndindex(x_shape[:first_axis]):
+                    for i, scalar_idxs in enumerate(zip(*vec_idxs)):  # noqa: B905
+                        out[(*outer, *scalar_idxs)] += y[(*outer, i)]
+                return out
 
         return advanced_inc_subtensor_multiple_vector