What causes "masked load_p" in jax-ml/jax?

Remove the mask: pad inputs so all loads are fully in-bounds and mask the arithmetic instead (e.g., zero out loaded values)

jax-ml/jax · error · NotImplementedError

masked load_p

Error message

masked load_p

What it means

The TPU Pallas interpreter has not implemented masked loads (primitives.load_p with a mask argument). Any pallas kernel using a masked pl.load raises NotImplementedError under interpret mode, regardless of whether the mask itself is valid.

Source

Thrown at jax/_src/pallas/mosaic/interpret/interpret_pallas_call.py:1342

      # array into the jaxpr when this function is traced.
      deferred_invals = functools.partial(env.read_many, eqn.invars)

      if (impl := _interpret_impls.get(prim, None)):
        invals = deferred_invals()
        # TODO(jburnim): Set up a proper kernel tracing environment for `impl`.
        impl_jaxpr = jax.make_jaxpr(functools.partial(impl, **eqn.params))(
            *invals)
        token, out = _interpret_jaxpr(
            impl_jaxpr, *impl_jaxpr.consts, *invals, ctx=ctx, token=token
        )
        if not prim.multiple_results:
          out = out[0]

      elif prim is primitives.load_p:
        (ref, transforms, mask, _) = jax.tree.unflatten(
            eqn.params['args_tree'], deferred_invals())
        if mask is not None:
          raise NotImplementedError('masked load_p')
        memory_space = _get_memory_space_and_raise_if_hbm(
            eqn.invars[0].aval, 'load_p'
        )
        ref, ref_transforms = mosaic_primitives._get_ref_and_transforms(ref)
        transforms = (*ref_transforms, *transforms)
        token, out = callback.io_callback(
            functools.partial(get, source_info=eqn.source_info),
            (TOKEN_SHAPE_DTYPE, eqn.outvars[0].aval),
            token,
            ctx.device_id,
            ctx.local_core_id,
            TPU_MEMORY_SPACE_IDXS[memory_space],
            ref,
            transforms,
        )

      elif prim is primitives.swap_p:
        (ref, transforms, val, mask) = jax.tree.unflatten(

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Remove the mask: pad inputs so all loads are fully in-bounds and mask the arithmetic instead (e.g., zero out loaded values)
Replace masked load with an unguarded load plus jnp.where on the loaded values for boundary semantics
Test on real TPU hardware/compile path if masked loads are required (interpret mode only limitation)

Example fix

# before
x = pl.load(ref, idx, mask=(rows < n), other=0.)
# after
x = pl.load(ref, idx)
x = jnp.where(rows[:, None] < n, x, 0.)

Defensive patterns

Strategy: fallback

Validate before calling

# pre-check kernel: pad inputs so masks are unnecessary
padded = jnp.pad(x, (0, (-x.shape[0]) % BM))  # all loads fully in-bounds

Try / catch

try:
    f_interpret(x)
except NotImplementedError as e:
    if 'masked load_p' in str(e):
        x_padded = jnp.pad(x, (0, (-x.shape[0]) % BM))
        f_interpret(x_padded)[:x.shape[0]]

Prevention

Avoid masked loads; pad inputs and mask arithmetic with jnp.where
Run interpret-mode smoke tests of every kernel primitive
Keep a hardware-run fallback for unsupported interpreter features

When it happens

Trigger: Calling pl.load(ref, indices, mask=...) inside a kernel executed with interpret=True (or the TPU mosaic interpreter path).

Common situations: Boundary-handling code with masked loads that works under GPU/Triton Pallas interpretation but is unsupported for TPU; testing TPU kernels locally in interpret mode.

Related errors

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/6740ef5c589bc256. Report an issue: GitHub.