jax-ml/jax · error · NotImplementedError

Scalar arguments not (yet) supported on GPU

Error message

Scalar arguments not (yet) supported on GPU

What it means

GPU interpret mode does not support scalar (index-operand) arguments to pallas kernels; if any scalar arguments are present, interpret_pallas_call raises NotImplementedError.

Source

Thrown at jax/_src/pallas/mosaic_gpu/interpret/interpret_pallas_call.py:420

      ordered=True,
  )

  token = gpu_callbacks.call_initialize_shared_memory(
      token=token,
      num_gpus=jnp.int32(device_info.num_devices),
      num_threads_per_block=jnp.int32(num_threads_per_block),
      num_blocks_per_cluster=jnp.int32(num_blocks_per_cluster),
      interpret_params=interpret_params,
  )

  dynamic_grid_args, scalars, inputs = split_list(
      args,
      [grid_mapping.num_dynamic_grid_bounds, grid_mapping.num_index_operands],
  )
  if dynamic_grid_args:
    raise NotImplementedError("Dynamic grid bounds not (yet) supported on GPU")
  if scalars:
    raise NotImplementedError("Scalar arguments not (yet) supported on GPU")

  assert grid_mapping.num_index_operands == 0

  token, input_buffer_keys = _allocate_buffers_for_inputs(
      token,
      device,
      jaxpr.invars[: grid_mapping.num_inputs],
      inputs,
  )

  token, output_buffers = _allocate_buffers_for_outputs(
      token,
      device,
      num_threads_per_block,
      input_output_aliases,
      grid_mapping,
      input_buffer_keys,
      inputs,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Remove scalar arguments; bake needed values into the kernel closure or Block transforms
2. Pass the values as 0-d arrays / buffers if the interpreter accepts them
3. Run on device where scalar operands may be supported
4. Update jax — GPU interpret coverage is expanding

Example fix

# before
kernel(x, start_idx)  # start_idx scalar operand
# after
start_idx = int(start_idx)
kernel = pallas_call(partial(fn, start=start_idx), out, grid=grid)
kernel(x)

assert not scalars, 'scalar args unsupported in GPU interpret mode'

When it happens

Trigger: 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.

Common situations: TPU pallas kernels using scalar index operands run against the GPU interpreter; passing python/jnp scalars where buffers are expected.

Related errors

• predicate not supported
• reduction_op not supported
• Dynamic grid bounds not (yet) supported in GPU interpret mod
• Dynamic grid bounds not (yet) supported on GPU
• Unsupported memory space: {space}

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/36ff12e74c7683b0. Report an issue: GitHub.