{"record":{"id":"fcaea41d31119889","repo":"jax-ml/jax","slug":"get-only-supports-loading-scalars-from-smem","errorCode":null,"errorMessage":"Get only supports loading scalars from SMEM.","messagePattern":"Get only supports loading scalars from SMEM\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/sc_lowering.py","lineNumber":135,"sourceCode":"        \" supported on SC\"\n    )\n  if not all(s == 1 for s in strides):\n    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],","sourceCodeStart":117,"sourceCodeEnd":153,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/sc_lowering.py#L117-L153","documentation":"On SC, scalar (0-d) loads are only supported from SMEM. Trying to load a scalar from VMEM (or another space) fails; instead load a 1-element array and extract the element afterwards.","triggerScenarios":"out_aval.ndim == 0 while ref_memory_space is not SMEM, e.g. v = vmem_ref[i, j] with scalar result inside an SC kernel.","commonSituations":"Scalar coefficient lookups from VMEM buffers; index math that collapses all dims.","solutions":["Load a length-1 dynamic slice and index element 0: v = ref[pl.ds(idx, 1)]; v[0]","Move the scalar into SMEM and read from there"],"exampleFix":"# before\ns = vmem_ref[i, j]\n# after\nv = vmem_ref[pl.ds(i, 1), pl.ds(j, 1)]\ns = v[0]","handlingStrategy":"fallback","validationCode":"if out_aval.ndim == 0 and space is not MemorySpace.SMEM:\n    v = ref[pl.ds(idx, 1)]  # load vector instead","typeGuard":"null","tryCatchPattern":"null","preventionTips":["Use pl.ds(idx, 1) + v[0] pattern for scalar loads from VMEM"],"tags":["jax","pallas","tpu","sparsecore","scalar-load"],"backgroundTag":"unsupported-scalar-memory-access","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}