{"record":{"id":"e4e3107b91fa8f19","repo":"jax-ml/jax","slug":"get-does-not-support-masked-scalar-loads","errorCode":null,"errorMessage":"Get does not support masked scalar loads","messagePattern":"Get does not support masked scalar loads","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/sc_lowering.py","lineNumber":138,"sourceCode":"    raise NotImplementedError(\n        \"Get only supports slices with stride 1, got {strides}\"\n    )\n\n  if (out_aval.ndim == 0) != (ref_memory_space is MemorySpace.SMEM):\n    message = \"Get only supports loading scalars from SMEM.\"\n    if ref_memory_space is MemorySpace.SMEM:\n      message += \" Trying to load an array of shape {out_aval.shape}.\"\n    elif ref_memory_space is MemorySpace.VMEM:\n      message += (\n          \" To load a scalar from VMEM, load an array first and then extract a\"\n          \" particular element, e.g. ``v = ref[pl.ds(idx, ...)]; v[0]``.\"\n      )\n    else:\n      message += f\" Trying to load a scalar from {ref_memory_space!r}.\"\n    raise NotImplementedError(message)\n  if out_aval.ndim == 0:\n    if mask is not None:\n      raise NotImplementedError(\"Get does not support masked scalar loads\")\n    return memref.load(ref, starts)\n\n  if not ctx.lowering_context.needs_layout_passes:\n    _check_aval_is_supported(\"Get\", out_aval)\n  out_vec_type = ir.VectorType.get(\n      out_aval.shape, _dtype_to_ir_type(out_aval.dtype)\n  )\n  if not ctx.lowering_context.needs_layout_passes:\n    return tpu.vector_load(\n        out_vec_type, ref, indices=starts, strides=[], mask=mask\n    )\n  # Load at the full memref rank, keeping integer-indexed dims as size 1,\n  # because apply-vector-layout requires the vector rank to match the memref.\n  memref_vec_shape = cast(\n      Sequence[int],\n      [1 if squeeze else s for s, squeeze in zip(sizes, squeeze_dims)],\n  )\n  memref_vec_type = ir.VectorType.get(","sourceCodeStart":120,"sourceCodeEnd":156,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/sc_lowering.py#L120-L156","documentation":"Masked scalar loads are not implemented in the SC lowering: when the load result is 0-d, passing a mask to Get raises immediately. Use an unmasked scalar load or handle bounds manually.","triggerScenarios":"v = ref[i, j] with mask=... (via pallas_get with a mask) where the result is a scalar, in a SparseCore kernel.","commonSituations":"Boundary-guarded lookups using masks, ported from TensorCore kernels.","solutions":["Clamp indices instead of masking for scalar loads","Load a 1-element masked vector then extract: v = masked_load(...); v[0] after ensuring valid lanes"],"exampleFix":"# before\nv = ref[i, j]  # with mask applied via pallas_get(mask=...)\n# after\ni = jnp.clip(i, 0, n - 1)\nv = ref[i, j]  # unmasked, index clamped","handlingStrategy":"fallback","validationCode":"if scalar and mask is not None: idx = jnp.clip(idx, 0, n - 1)","typeGuard":"null","tryCatchPattern":"null","preventionTips":["Clamp indices instead of masking scalar loads"],"tags":["jax","pallas","tpu","sparsecore","masked-load"],"backgroundTag":"unsupported-masked-operation","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}