JAX v0.11 compatibility: `scan` param change breaks detection, new `empty` and `call_hi_primitive` primitives unhandled

[adrhill]

JAX v0.11.0 (release notes, published 2026-07-16) breaks sparsity detection in two places and adds one forward-looking gap.
Everything below was verified empirically against jax==0.11.0 (Python 3.14.2) by running the full test suite and probing jaxprs directly.

Test suite under jax==0.11.0: 21 failed, 3982 passed. All 21 failures are in tests/_interpret/test_scan.py with KeyError: 'num_consts'.

1. scan params changed: num_consts/num_carry replaced by ft_in/ft_out

Not mentioned in the release notes, found only by running the test suite.
The scan primitive params changed:

# jax 0.10.2
sorted(eqn.params.keys())
# ['jaxpr', 'length', 'num_carry', 'num_consts', 'reverse', 'unroll']

# jax 0.11.0
sorted(eqn.params.keys())
# ['ft_in', 'ft_out', 'jaxpr', 'length', 'reverse', 'unroll']

_prop_scan in _interpret/_scan.py reads eqn.params["num_consts"] and eqn.params["num_carry"] to split invars into consts, carry, and xs, so every scan-containing function now raises KeyError: 'num_consts'.

The new params are jax._src.flattree.FTTuple objects grouping the flat invars and outvars by role:

# scan with 1 closure const, 2 carries, 1 xs (4 invars, 3 outvars):
eqn.params["ft_in"]   # ((None,), (None, None), (None,))  = (consts, carry, xs)
eqn.params["ft_out"]  # ((None, None), (RightsOnly(Right(None)),))  = (carry, ys)

Needed change

Recover the counts from the group sizes, keeping the old path for jax < 0.11:

if "num_consts" in eqn.params:
    num_consts = eqn.params["num_consts"]
    num_carry = eqn.params["num_carry"]
else:
    num_consts = len(eqn.params["ft_in"].elts[0])
    num_carry = len(eqn.params["ft_in"].elts[1])

Group sizes were verified to match the flat invar/outvar layout for tuple carries.
The fix should double-check the pytree cases against the existing test_scan_pytree_xs / test_scan_pytree_ys tests, since FTTuple elements can in principle carry tree structure.
The linear param is also gone in 0.11, which only affects the docstring param list in _scan.py (the handler never read it).

2. New empty primitive from jnp.empty / jnp.empty_like

JAX 0.11.0 changed jax.numpy.empty and jax.numpy.empty_like to produce genuinely uninitialized arrays.
They used to lower to zeros:

# jax 0.10.2
jax.make_jaxpr(lambda: jnp.empty((3,)))()
# { lambda ; . let a:f32[3] = broadcast_in_dim 0.0:f32[] in (a,) }

# jax 0.11.0
jax.make_jaxpr(lambda: jnp.empty((3,)))()
# { lambda ; . let a:f32[3] = empty[dtype=float32 out_sharding=None shape=(3,)] in (a,) }

Detection on any function using jnp.empty now raises:

def f(x):
    out = jnp.empty((3,))
    return out.at[:].set(x[:3] * 2.0)

asdex.jacobian_sparsity(f, jnp.ones(4))
# NotImplementedError: No handler for primitive 'empty'. ...

No asdex test currently exercises jnp.empty, which is why the suite did not catch this.

Needed change

Add a _prop_empty handler (new _empty.py module).
The primitive is nullary, so the handler mirrors _prop_iota without the const tracking:

• state_indices[eqn.outvars[0]] = _empty_index_sets(numel) with numel from eqn.params["shape"].
• Do not record anything in state_consts.
  The values are uninitialized, unlike the old zeros lowering, so zero-clearing in mul and zero-skipping in dot_general must not fire.
  Functions that relied on jnp.empty returning zeros may therefore detect denser patterns than under JAX 0.10, which is correct.
• Add a detection test covering jnp.empty followed by .at[...].set(...).

3. New call_hi_primitive from hijax-based remat (jax_remat3, jax.custom_remat)

JAX 0.11.0 adds jax.custom_remat and the jax_remat3 implementation.
With jax.config.update("jax_remat3", True), jax.checkpoint no longer emits remat2 (handled via _prop_closed_jaxpr) but the generic hijax call primitive:

jax.config.update("jax_remat3", True)
jax.make_jaxpr(jax.checkpoint(lambda y: jnp.sin(y) * y))(jnp.ones(3))
# { lambda ; a:f32[3]. let
#     b:f32[3] = call_hi_primitive[
#       _prim=RematTraced[{'jaxpr': { lambda ; c:f32[3]. ... }, 'policy': None}]
#     ] a
#   in (b,) }

jax.custom_remat emits the same primitive with _prim=CustomRemat[{'jaxpr': ..., ...}].
The flag is off by default in 0.11.0, so plain jax.checkpoint still works with asdex today.
But custom_remat requires jax_remat3, and the flag looks like the future default.

The inner jaxpr is not in eqn.params["jaxpr"] but nested inside the hi-primitive object:

prim = eqn.params["_prim"]       # RematTraced, a VJPHiPrimitive subclass
type(prim.jaxpr)                 # jax._src.core.Jaxpr (open jaxpr, same as remat2's param)
prim.params                      # {'jaxpr': ..., 'policy': None}

Needed change

Add a call_hi_primitive case to the dispatch in _interpret/__init__.py:

• If eqn.params["_prim"] has a jaxpr, recurse into it the same way _prop_closed_jaxpr does (forward consts and bounds from eqn.invars to the inner invars).
  The existing helper almost fits, it just needs a variant that takes the jaxpr directly instead of a param_key.
• If the hi-primitive is opaque (a user-defined HiPrimitive without a wrapped jaxpr), fall through to _prop_throw_error so we never silently guess.

Not affected

• The jax.core / jax.interpreters.pxla API removals: asdex only imports jax.core.Tracer (in tests/test_pattern.py), which was not removed.
• while and cond params are unchanged, all their tests pass under 0.11.0.
• All .empty( calls in asdex itself are np.empty (NumPy), not jnp.empty.
• Dropped Python 3.11 / NumPy 2.0 / SciPy 1.14 support: no pin changes needed for jax>=0.9.0 compatibility, older JAX versions keep working on the existing matrix.

[gdalle]

Oops I opened #180 before I saw this one