jax-ml/jax · error · ValueError

Unsupported transform: {transform}

Error message

Unsupported transform: {transform}

What it means

Internal validation in the GPU interpreter: only NDIndexer transforms are understood when applying a Block's transforms to host buffers. Any other transform type reaching _validate_transforms is rejected.

Source

Thrown at jax/_src/pallas/mosaic_gpu/interpret/gpu_callbacks.py:517

def _is_dynamic(indexer: indexing.NDIndexer) -> bool:
  return any(
      isinstance(idx, indexing.Slice)
      and (idx.is_dynamic_start or idx.is_dynamic_size)
      for idx in indexer.indices
  )


def _validate_transforms(transforms):
  for transform in transforms:
    match transform:
      case indexing.NDIndexer():
        if _is_dynamic(transform):
          raise ValueError(
              "Dynamic indexing not supported in GPU interpret mode"
          )
      case _:
        raise ValueError(f"Unsupported transform: {transform}")


def _get(
    token: jax.Array,
    mesh_location: memory.MeshLocation,
    thread: memory.Thread | None,
    allocation_key_as_array: jax.Array,
    transforms,
    block_indices=None,
    grid_loop_idx=None,
    clock=None,
    increment_clock: bool = True,
    source_info=None,
    input_name=None,
) -> tuple[jax.Array, np.ndarray]:
  """Performs a read from the buffer for `allocation_key_as_array` from the given device and thread."""
  allocation_key = HostAllocationKey.from_array(allocation_key_as_array)
  del allocation_key_as_array

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Check the repr in the message to identify the transform and avoid constructing Blocks that produce it
2. Update jax (and the mosaic plugin) to matching versions where the transform is supported
3. If it's an NDIndexer wrapped differently, ensure indices are plain ints/static so it validates as NDIndexer
4. Report upstream with a minimal repro if it's a legitimate transform

try:
    kernel(x)
except ValueError as e:
    if 'Unsupported transform' in str(e):
        simplify Block construction (plain ds[...] indexing)

When it happens

Trigger: Rare; occurs if a Block reference carries a transform not representable as indexing.NDIndexer (e.g. a new/unsupported transform kind introduced by API changes), passed through _get or _swap in interpret mode.

Common situations: Version skew between jax and jax_plugins/pallas code paths; using experimental Block/transform features not yet implemented in the interpreter.

Related errors

• `thread_id` must be zero when allocating a buffer for all th
• `block_id` must be zero when allocating a buffer for all thr
• Out-of-bounds read of {allocation_key}: reading [{read_range
• Out-of-bounds block index {block_indices} for {allocation_ke
• Dynamic indexing not supported in GPU interpret mode

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