{"record":{"id":"b15375dcd1e6fd03","repo":"jax-ml/jax","slug":"indexing-into-a-shaped-ref-not-yet-supported-on","errorCode":null,"errorMessage":"Indexing into a ()-shaped Ref not yet supported on TPU.","messagePattern":"Indexing into a \\(\\)-shaped Ref not yet supported on TPU\\.","errorType":"validation","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/lowering.py","lineNumber":2327,"sourceCode":"    idx = cast(NDIndexer, idx)\n    if idx.int_indexer_shape:\n      raise NotImplementedError()\n    elt_slices = [\n        indexing.Slice(0, size) for size in physical_element_aval.shape]\n    idx = NDIndexer(\n        indices=idx.indices + tuple(elt_slices),\n        shape=idx.shape + physical_element_aval.shape,\n        int_indexer_shape=(),\n    )\n    physical_out_dtype = physical_element_aval.dtype\n    physical_out_shape = jax_core.physical_shape(\n        aval_out.shape, aval_out.dtype\n    )\n  else:\n    physical_out_dtype = aval_out.dtype\n    physical_out_shape = aval_out.shape\n  if not is_smem_load and not ref_block_shape:\n    raise NotImplementedError(\n        \"Indexing into a ()-shaped Ref not yet supported on TPU.\")\n  starts, sizes, strides, _, _ = _indexer_to_start_size_stride(\n      idx,\n      ref_block_shape,\n      cast_to_index=True,\n  )\n  need_stride = not all((s is None or s == 1) for s in strides)\n  if is_smem_load:\n    if ctx.avals_out[0].shape:\n      raise ValueError(\"Can only load scalars from SMEM\")\n    return _maybe_cast_load_to_bool(ctx, aval_out, memref.load(ref, starts))\n  elif str(ref_type.memory_space) != \"#tpu.memory_space<vmem>\":\n    extra = \"\"\n    if str(ref_type.memory_space) == \"#tpu.memory_space<any>\":\n      extra = \" ANY memory space can only be accessed using async_copy.\"\n    raise ValueError(\n        \"Loads are only allowed on VMEM and SMEM references.\" + extra\n    )","sourceCodeStart":2309,"sourceCodeEnd":2345,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/lowering.py#L2309-L2345","documentation":"Raised when a Pallas TPU load targets a ()-shaped (scalar) Ref that is not in SMEM. VMEM ops on TPU are vector ops and cannot address a scalar block, so scalar loads are only supported from SMEM; loading a scalar VMEM block is not yet implemented.","triggerScenarios":"pl.load on a Ref whose block shape is () (scalar block) while the ref lives in VMEM (not SMEM). Typically from a BlockSpec with scalar block shape for a scalar input.","commonSituations":"Passing scalars (like a loop counter or scalar hyperparameter) into the kernel with default VMEM BlockSpec; squeezing a dimension away so the block becomes scalar.","solutions":["Use memory_space=pltpu.MemorySpace.SMEM for the scalar input's BlockSpec","Or give the value shape (1,) instead of () and keep it in VMEM","Pass scalars as compile-time constants or via inlined constants rather than refs"],"exampleFix":"# before\nspec = pltpu.BlockSpec(())\n# after\nspec = pltpu.BlockSpec((), memory_space=pltpu.MemorySpace.SMEM)","handlingStrategy":"validation","validationCode":"def scalar_spec_ok(spec):\n    block = getattr(spec, 'block_shape', None)\n    ms = str(getattr(spec, 'memory_space', '') or '')\n    if block == () and 'SMEM' not in ms:\n        raise ValueError('scalar block must use memory_space=SMEM on TPU')","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Prefer (1,) block shapes over () for VMEM values","Route all scalar inputs through SMEM specs via a shared helper"],"tags":["jax","pallas","tpu","scalar-ref","smem"],"backgroundTag":"pallas-scalar-ref-unsupported","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}