Fix JAX AdvancedIncSubtensor crashing on a 0d update

set_subtensor(x[idx_vector], 1.0) raises AttributeError: 'float' object has
no attribute 'shape' under the JAX backend. The update input is a 0d
TensorType, so _check_runtime_broadcast_of_vector_index is entitled to a
value with a .shape, but jax_typify downgrades 0d arrays to Python scalars
before the op ever sees them.

Coerce y back to an array at the dispatch, which is the layer that knows it
is handing linker values to a check written against arrays. The .item() in
jax_typify is load-bearing for dispatches that need a static Python scalar
(axis/shape/loop-bound arguments), so it is left alone.

Author: velochy

pytensor/link/jax/dispatch/subtensor.py

@@ -1,3 +1,5 @@
+import jax.numpy as jnp
+
 from pytensor.link.jax.dispatch.basic import jax_funcify
 from pytensor.tensor.subtensor import (
     AdvancedIncSubtensor,
@@ -60,7 +62,11 @@ def incsubtensor(x, y, *ilist, jax_fn=jax_fn, idx_list=op.idx_list):
             indices = indices[0]
 
         if isinstance(op, AdvancedIncSubtensor):
-            op._check_runtime_broadcast_of_vector_index(node, x, y, indices)
+            # jax_typify downgrades 0d arrays to Python scalars, which have no .shape,
+            # so re-arrayify y for a check that is written against array values.
+            op._check_runtime_broadcast_of_vector_index(
+                node, x, jnp.asarray(y), indices
+            )
 
         return jax_fn(x, indices, y)
 

tests/link/jax/test_subtensor.py

@@ -225,6 +225,26 @@ def test_jax_IncSubtensor():
     compare_jax_and_py([], [out_pt], [])
 
 
+@pytest.mark.parametrize(
+    "func", (pt_subtensor.advanced_inc_subtensor1, pt_subtensor.advanced_set_subtensor1)
+)
+def test_jax_AdvancedIncSubtensor1_scalar_y(func):
+    """A 0d update must not trip the runtime-broadcast check.
+
+    The check is written against array values, but ``jax_typify`` turns a 0d array into a
+    Python scalar, which has no ``.shape``.
+    """
+    from pytensor import function
+
+    x = pt.zeros((5,))
+    out = func(x, 1.0, np.array([1, 3]))
+
+    res = function([], out, mode="JAX")()
+    expected = np.zeros(5)
+    expected[[1, 3]] = 1.0
+    np.testing.assert_allclose(res, expected)
+
+
 @pytest.mark.parametrize(
     "func", (pt_subtensor.advanced_inc_subtensor1, pt_subtensor.advanced_set_subtensor1)
 )