{"record":{"id":"36ff12e74c7683b0","repo":"jax-ml/jax","slug":"scalar-arguments-not-yet-supported-on-gpu","errorCode":null,"errorMessage":"Scalar arguments not (yet) supported on GPU","messagePattern":"Scalar arguments not \\(yet\\) supported on GPU","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/interpret/interpret_pallas_call.py","lineNumber":420,"sourceCode":"      ordered=True,\n  )\n\n  token = gpu_callbacks.call_initialize_shared_memory(\n      token=token,\n      num_gpus=jnp.int32(device_info.num_devices),\n      num_threads_per_block=jnp.int32(num_threads_per_block),\n      num_blocks_per_cluster=jnp.int32(num_blocks_per_cluster),\n      interpret_params=interpret_params,\n  )\n\n  dynamic_grid_args, scalars, inputs = split_list(\n      args,\n      [grid_mapping.num_dynamic_grid_bounds, grid_mapping.num_index_operands],\n  )\n  if dynamic_grid_args:\n    raise NotImplementedError(\"Dynamic grid bounds not (yet) supported on GPU\")\n  if scalars:\n    raise NotImplementedError(\"Scalar arguments not (yet) supported on GPU\")\n\n  assert grid_mapping.num_index_operands == 0\n\n  token, input_buffer_keys = _allocate_buffers_for_inputs(\n      token,\n      device,\n      jaxpr.invars[: grid_mapping.num_inputs],\n      inputs,\n  )\n\n  token, output_buffers = _allocate_buffers_for_outputs(\n      token,\n      device,\n      num_threads_per_block,\n      input_output_aliases,\n      grid_mapping,\n      input_buffer_keys,\n      inputs,","sourceCodeStart":402,"sourceCodeEnd":438,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/interpret/interpret_pallas_call.py#L402-L438","documentation":"GPU interpret mode does not support scalar (index-operand) arguments to pallas kernels; if any scalar arguments are present, interpret_pallas_call raises NotImplementedError.","triggerScenarios":"Calling a pallas kernel that takes scalar operands (values in the scalars/index-operands slot of the grid mapping) under GPU interpret mode — e.g. kernels whose blocks take integer scalars computed at launch.","commonSituations":"TPU pallas kernels using scalar index operands run against the GPU interpreter; passing python/jnp scalars where buffers are expected.","solutions":["Remove scalar arguments; bake needed values into the kernel closure or Block transforms","Pass the values as 0-d arrays / buffers if the interpreter accepts them","Run on device where scalar operands may be supported","Update jax — GPU interpret coverage is expanding"],"exampleFix":"# before\nkernel(x, start_idx)  # start_idx scalar operand\n# after\nstart_idx = int(start_idx)\nkernel = pallas_call(partial(fn, start=start_idx), out, grid=grid)\nkernel(x)","handlingStrategy":"validation","validationCode":"assert not scalars, 'scalar args unsupported in GPU interpret mode'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Bake scalars into the kernel closure via functools.partial","Pass data as arrays, not scalar operands"],"tags":["pallas","mosaic-gpu","interpret-mode","scalar-args","not-implemented"],"backgroundTag":"unsupported-operation-in-interpreter","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}