{"record":{"id":"6740ef5c589bc256","repo":"jax-ml/jax","slug":"masked-load-p","errorCode":null,"errorMessage":"masked load_p","messagePattern":"masked load_p","errorType":"validation","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/interpret/interpret_pallas_call.py","lineNumber":1342,"sourceCode":"      # array into the jaxpr when this function is traced.\n      deferred_invals = functools.partial(env.read_many, eqn.invars)\n\n      if (impl := _interpret_impls.get(prim, None)):\n        invals = deferred_invals()\n        # TODO(jburnim): Set up a proper kernel tracing environment for `impl`.\n        impl_jaxpr = jax.make_jaxpr(functools.partial(impl, **eqn.params))(\n            *invals)\n        token, out = _interpret_jaxpr(\n            impl_jaxpr, *impl_jaxpr.consts, *invals, ctx=ctx, token=token\n        )\n        if not prim.multiple_results:\n          out = out[0]\n\n      elif prim is primitives.load_p:\n        (ref, transforms, mask, _) = jax.tree.unflatten(\n            eqn.params['args_tree'], deferred_invals())\n        if mask is not None:\n          raise NotImplementedError('masked load_p')\n        memory_space = _get_memory_space_and_raise_if_hbm(\n            eqn.invars[0].aval, 'load_p'\n        )\n        ref, ref_transforms = mosaic_primitives._get_ref_and_transforms(ref)\n        transforms = (*ref_transforms, *transforms)\n        token, out = callback.io_callback(\n            functools.partial(get, source_info=eqn.source_info),\n            (TOKEN_SHAPE_DTYPE, eqn.outvars[0].aval),\n            token,\n            ctx.device_id,\n            ctx.local_core_id,\n            TPU_MEMORY_SPACE_IDXS[memory_space],\n            ref,\n            transforms,\n        )\n\n      elif prim is primitives.swap_p:\n        (ref, transforms, val, mask) = jax.tree.unflatten(","sourceCodeStart":1324,"sourceCodeEnd":1360,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/interpret/interpret_pallas_call.py#L1324-L1360","documentation":"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.","triggerScenarios":"Calling pl.load(ref, indices, mask=...) inside a kernel executed with interpret=True (or the TPU mosaic interpreter path).","commonSituations":"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.","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)"],"exampleFix":"# before\nx = pl.load(ref, idx, mask=(rows < n), other=0.)\n# after\nx = pl.load(ref, idx)\nx = jnp.where(rows[:, None] < n, x, 0.)","handlingStrategy":"fallback","validationCode":"# pre-check kernel: pad inputs so masks are unnecessary\npadded = jnp.pad(x, (0, (-x.shape[0]) % BM))  # all loads fully in-bounds","typeGuard":null,"tryCatchPattern":"try:\n    f_interpret(x)\nexcept NotImplementedError as e:\n    if 'masked load_p' in str(e):\n        x_padded = jnp.pad(x, (0, (-x.shape[0]) % BM))\n        f_interpret(x_padded)[:x.shape[0]]","preventionTips":["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"],"tags":["jax","pallas","tpu","interpret-mode","not-implemented","masked-load"],"backgroundTag":"unsupported-feature","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}