Fix AdvancedIncSubtensor runtime-broadcast check crashing on a scalar update

[velochy]

What / why

set_subtensor(x[idx_vector], scalar) works on the Python and numba backends and raises on JAX:

import numpy as np, pytensor, pytensor.tensor as pt

x = pt.vector("x")
out = pt.set_subtensor(x[pt.constant([1, 3])], 1.0)   # 0d update

pytensor.function([x], out, mode="JAX")(np.zeros(5))
# AttributeError: 'float' object has no attribute 'shape'
#   pytensor/link/jax/dispatch/subtensor.py:63  in incsubtensor
#   pytensor/tensor/subtensor.py:2468           in _check_runtime_broadcast_of_vector_index

The update input is a 0d TensorType, so _check_runtime_broadcast_of_vector_index is entitled to
a value with a .shape. It doesn't get one because jax_typify downgrades 0d arrays to Python
scalars:

@jax_typify.register(np.ndarray)
def jax_typify_ndarray(data, dtype=None, **kwargs):
    if len(data.shape) == 0:
        return data.item()          # 0d array -> Python scalar
    return jnp.array(data, dtype=dtype)

The change

Coerce at the dispatch call site — the layer that knows it is handing linker values to a check
written against arrays:

-op._check_runtime_broadcast_of_vector_index(node, x, y, indices)
+op._check_runtime_broadcast_of_vector_index(node, x, jnp.asarray(y), indices)

This is a bandaid over the typify behaviour, by @ricardoV94's call in the thread below. Removing the
.item() branch instead does fix the crash on its own, but it takes tests/link/jax/ from 2
failures to 27 (25 new, verified against the same suite on clean main) — 10 in test_scan.py,
4 in test_sort.py (ValueError: Non-hashable static arguments), 3 in test_shape.py, plus
linalg/pad/conv, mostly TracerBoolConversionError and non-hashable static args. The
.item() is load-bearing for dispatches that want a static Python value for an axis, shape or loop
bound, so doing it properly means having those pull their scalar from the constant node instead —
a much wider change than this PR.

Testing

• New test_jax_AdvancedIncSubtensor1_scalar_y, parametrised over
  advanced_inc_subtensor1 / advanced_set_subtensor1 to match the neighbouring
  test_jax_AdvancedIncSubtensor1_runtime_broadcast. Verified it fails without the fix
  (AttributeError) and passes with it.
• The existing test_jax_AdvancedIncSubtensor1_runtime_broadcast still passes: a length-1 y
  written into 2 indices continues to raise ValueError: Runtime broadcasting not allowed…, so the
  guard still fires where it should.
• tests/link/jax/test_subtensor.py passes (11 passed, 1 xfailed). tests/tensor/test_subtensor.py
  has one failure, TestSubtensor::test_boolean, which fails identically on the base commit.
• Found via real models: PyMC ranking models (Multinomial over best/worst choice sets) whose
  expansion function builds its rank structure with a scalar set_subtensor. Under
  nuts_sampler="nutpie" with the JAX backend this surfaced only as an opaque
  Error during point expansion panic from nuts-rs mid-sampling, with the Python exception
  swallowed. Three multi-step MRP models (~5k respondents) that could not be fitted under the JAX
  backend at all now complete; the numba backend fit all three identically before and after.

Versions: pytensor main @ a85bc78, jax 0.6.2, numpy 2.3.5, Python 3.12, Linux.

Related: #2241 extended this check to the MLX backend and #2303 fixed a different false positive in
it; neither covers the 0d-y case. The MLX and PyTorch dispatches call the same method, so they are
presumably reachable the same way if their typify does the same thing — not touched here, since I
have not reproduced it on either.

[velochy]

Separate from this fix, while verifying that the guard still fires I found that the numba backend
does not enforce this check at all.

Same graph, a length-1 y written into 2 indices, each linker pinned explicitly:

import numpy as np, pytensor, pytensor.tensor as pt
from pytensor.compile.mode import Mode

x, y = pt.vector("x"), pt.vector("y")
out = pt.set_subtensor(x[pt.constant([1, 3])], y)

for label, mode in [
    ("py    ", Mode(linker="py",    optimizer=None)),
    ("cvm/C ", Mode(linker="cvm",   optimizer="fast_run")),
    ("numba ", Mode(linker="numba", optimizer="fast_run")),
    ("jax   ", "JAX"),
]:
    try:
        print(label, pytensor.function([x, y], out, mode=mode)(np.zeros(5), np.ones(1)))
    except Exception as e:
        print(label, f"{type(e).__name__}: {str(e).splitlines()[0]}")

py     ValueError: Runtime broadcasting not allowed. ...
cvm/C  ValueError: Runtime broadcasting not allowed. ...
numba  [0. 1. 0. 1. 0.]
jax    ValueError: Runtime broadcasting not allowed. ...

So perform, the C/VM linker and JAX all reject it, and (by call site) MLX and PyTorch would too —
numba is the only path that silently broadcasts.

Worth being explicit that this is not a missing call that could be dropped in next to the others.
The numba dispatch broadcasts deliberately as part of its codegen —
np.broadcast_to(y_adv_dims_front, (*adv_idx_shape, *basic_idx_shape)) in
link/numba/dispatch/subtensor.py — and the shapes are only known at run time, so matching the
other backends means emitting a shape comparison and a raise inside the generated hot-loop code.
That has a runtime cost, and it would make graphs that work today start raising for anyone relying
on the current numba behaviour.

Given that, it seemed like your call rather than something to fold into a crash fix. Happy to open
it as its own issue, or to put up a PR if you would like numba brought in line.

One caveat on my own numbers: FAST_RUN resolves to the numba linker in my environment
(config.linker = auto), so if you reproduce this with mode="FAST_RUN" on a box configured for
the C backend you will see the C behaviour (raises), not the numba one. The explicit Mode(...)
above avoids that ambiguity.

[velochy]

@ricardoV94 I know I am maxing out my PR quota again, but this one seems small and contained... hopefully :)

[ricardoV94]

May sound pedantic but that AdvancedIncSubtensor is expecting a 0d array type, not a scalar, so the bug is actually elsewhere, a loose typify or another op.

>>> import numpy as np, pytensor, pytensor.tensor as pt
... 
... x = pt.vector("x")
... out = pt.set_subtensor(x[pt.constant([1, 3])], 1.0)   # scalar update
... out.dprint(print_type=True)
... 
AdvancedSetSubtensor [id A] <Vector(float64, shape=(?,))>
 ├─ x [id B] <Vector(float64, shape=(?,))>
 ├─ 1.0 [id C] <Scalar(float32, shape=())>
 └─ [1 3] [id D] <Vector(int64, shape=(2,))>

That scalar is a 0d float32 array. Should have a shape. But JAX typify defaults to "everything is fine", that's where the bug is.

[velochy]

You're right, and it isn't pedantic — I was patching the symptom. The input is a 0d TensorType,
so the check is entitled to a value with a .shape; the reason it doesn't get one is here:

@jax_typify.register(np.ndarray)
def jax_typify_ndarray(data, dtype=None, **kwargs):
    if len(data.shape) == 0:
        return data.item()          # 0d array -> Python scalar
    return jnp.array(data, dtype=dtype)

jax_typify downgrades every 0d array to a Python scalar, so the .shape is gone before the op
ever sees it. Fixing that instead of the check does resolve the crash on its own, with
subtensor.py untouched.

But dropping the .item() branch outright is not viable as-is — tests/link/jax/ goes from
2 failures to 27 (25 new, verified against the same suite on clean main):

	failures
clean main	2
.item() branch removed	27

The new ones cluster where JAX needs a static Python value rather than an array — 10 in
test_scan.py, 4 in test_sort.py (ValueError: Non-hashable static arguments), 3 in
test_shape.py, plus linalg/pad/conv — mostly TracerBoolConversionError and non-hashable
static args. So the .item() looks load-bearing for axis/shape/loop-bound arguments, and the real
fix is presumably to keep 0d values as arrays at typify and have the dispatches that need a static
scalar pull it from the constant node instead. That is a much wider change than this PR, and enough
of a design call that I would rather not guess at it.

How would you like to play it? Options as I see them:

1. I close this and open an issue describing the typify bug with the failure breakdown above, for
   someone to do properly.
2. I take on the typify change — happy to, but it needs a decision on how the static-argument
   dispatches should get their scalars, and it will touch a fair number of files.
3. Narrow the current PR to coercing at the JAX dispatch call site
   (op._check_runtime_broadcast_of_vector_index(node, x, jnp.asarray(y), indices)) — still a
   workaround, but at the layer that knows it is handing values to a check that expects arrays,
   rather than loosening the check itself.

Happy with whichever; I have the reproducer and the suite runs set up either way. For context on why
I chased it at all: it makes any PyMC model with a ranking likelihood unsamplable under the JAX
backend (nutpie surfaces it only as an opaque Error during point expansion from nuts-rs), so it is
not purely theoretical — but that is an argument for fixing it right rather than fast.

[ricardoV94]

3 sounds the best bandaid

[velochy]

Done — narrowed to option 3. subtensor.py is untouched now; the whole change is jnp.asarray(y) at the JAX dispatch call site, plus the regression test (verified failing without it). PR description updated to describe the typify root cause rather than the check.

Happy to open a follow-up issue for the typify part if you want it tracked.