{"record":{"id":"823c5a9584b28651","repo":"jax-ml/jax","slug":"aliasing-of-scalar-prefetch-arguments-is-not-curre","errorCode":null,"errorMessage":"Aliasing of scalar prefetch arguments is not currently supported in TPU interpret mode.","messagePattern":"Aliasing of scalar prefetch arguments is not currently supported in TPU interpret mode\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/interpret/interpret_pallas_call.py","lineNumber":2055,"sourceCode":"  input_buffer_ids = []\n  for i, var in enumerate(\n      jaxpr.invars[grid_mapping.num_index_operands:][:grid_mapping.num_inputs]):\n    assert var.aval.dtype == input_args[i].dtype  # pyrefly: ignore[missing-attribute]\n    token, buffer_id = callback.io_callback(\n        _allocate_buffer,\n        (TOKEN_SHAPE_DTYPE, jax.ShapeDtypeStruct((), jnp.int16)),\n        token,\n        device_id,\n        None,  # local_core_id\n        TPU_MEMORY_SPACE_IDXS[mosaic_core.MemorySpace.HBM],\n        input_args[i],\n    )\n    input_buffer_ids.append(buffer_id)\n\n  # Allocate buffers in HBM for pallas_call outputs.\n  oi_alias_map = {v: k - len(scalars) for k, v in input_output_aliases}\n  if any(i < 0 for i in oi_alias_map.keys()):\n    raise ValueError('Aliasing of scalar prefetch arguments is not currently '\n                     'supported in TPU interpret mode.')\n  output_buffer_ids = []\n  output_buffer_shapes = []\n  output_vals = []\n  num_outputs = grid_mapping.num_outputs\n  output_block_shapes = block_shapes[num_inputs : num_inputs + num_outputs]\n  for i, bm in enumerate(grid_mapping.block_mappings_output):\n    if i in oi_alias_map:\n      # Reuse the HBM buffer for the aliased pallas_call input.\n      output_buffer_ids.append(input_buffer_ids[oi_alias_map[i]])\n      output_buffer_shapes.append(input_args[oi_alias_map[i]].shape)\n      output_vals.append(input_args[oi_alias_map[i]])\n    else:\n      out_val = interpret_utils.get_uninitialized_array(\n          bm.array_aval.shape, bm.array_aval.dtype,\n          interpret_params.uninitialized_memory)\n      padded_val = interpret_utils.pad_to_block_dimension(\n          out_val, output_block_shapes[i], interpret_params.uninitialized_memory","sourceCodeStart":2037,"sourceCodeEnd":2073,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/interpret/interpret_pallas_call.py#L2037-L2073","documentation":"In TPU interpret mode, pallas_call input_output_aliases that alias a scalar prefetch argument (an input index below the number of scalar inputs) are rejected, because the interpreter's buffer allocation scheme cannot map scalar prefetch inputs to output buffer ids.","triggerScenarios":"Using input_output_aliases=(i, j) where i < len(scalars) — i.e., aliasing an output to a scalar (prefetch) input — while running with interpret mode on TPU.","commonSituations":"Kernels designed for GPU/compiled TPU that return an updated scalar argument (e.g., running token/seed) via output aliasing; porting such kernels to interpret-mode debugging.","solutions":["Alias outputs only to non-scalar (array) inputs","Return updated scalars as separate kernel outputs instead of aliasing the scalar input","Skip interpret mode for kernels requiring scalar aliasing"],"exampleFix":"# before\npallas_call(kernel, out_shapes, inputs, input_output_aliases=(0, 0))  # input 0 is scalar\n# after\npallas_call(kernel, out_shapes + scalar_shape, inputs)  # return updated scalar as extra output","handlingStrategy":"validation","validationCode":"num_scalars = sum(1 for x in inputs if getattr(x, 'shape', None) in ((), None))\nfor i, _ in input_output_aliases:\n    assert i >= num_scalars, 'cannot alias scalar prefetch inputs in TPU interpret mode'","typeGuard":null,"tryCatchPattern":"try:\n    interpret_run(kernel)\nexcept ValueError as e:\n    if 'scalar prefetch' in str(e):\n        # move aliased scalar to an explicit output and rerun\n        raise","preventionTips":["Alias outputs only to array inputs","Return updated scalars as extra outputs","Check alias indices against scalar input count before interpret runs"],"tags":["jax","pallas","tpu","aliasing","scalar-prefetch","interpret-mode"],"backgroundTag":"unsupported-feature","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}