jax-ml/jax · error · NotImplementedError

get not supported yet

Error message

get not supported yet

What it means

Raised by _get_pull_rule in the Pallas fuser when tracing a `get` on a Ref whose tree of indexers contains more than one NDIndexer. The block-spec propagation logic only handles a single indexer per get, so a pytree of multiple indexers (e.g. getting from multiple refs/indexers via one primitive) is rejected.

Source

Thrown at jax/_src/pallas/fuser/block_spec.py:1619

      _slice(i, b)
      for i, b in zip(block_idx, block_spec.block_shape, strict=True)
  )
  return ref.swap(val, idx=indexer)


@register_pull_block_spec_rule(state_primitives.get_p)
def _get_pull_rule(
    ctx: PullRuleContext, block_transform: BlockIndexTransform, *, tree
):
  if block_transform.block_shape is None:
    return [block_transform] + [no_block_index_transform] * (
        len(ctx.avals_in) - 1
    )
  ref_aval = ctx.avals_in[0]
  assert hasattr(ref_aval, 'shape')
  indexers_avals = tree_util.tree_unflatten(tree, ctx.avals_in[1:])
  if len(indexers_avals) > 1:
    raise NotImplementedError('get not supported yet')
  if not indexers_avals:
    indexer_aval = indexing.NDIndexer.make_trivial_indexer(ref_aval.shape)
  else:
    indexer_aval = indexers_avals[0]
  block_shape_iter = iter(block_transform.block_shape)
  block_shape = []
  if not all(
      bd is None
      or isinstance(bd, (int, pallas_core.Blocked, pallas_core.Squeezed))
      for bd in block_transform.block_shape
  ):
    raise NotImplementedError('get not supported yet')
  for idx_aval, size in zip(indexer_aval.indices, ref_aval.shape, strict=True):
    if not isinstance(idx_aval, indexing.Slice):
      assert hasattr(idx_aval, 'shape') and not idx_aval.shape
      block_shape.append(pallas_core.Squeezed())
      continue
    if not isinstance(idx_aval.start, int):

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use a single indexer per get: split multi-indexer gets into separate ref.get(...) calls
2. Flatten any pytree indexers into one NDIndexer before calling get
3. Check for accidental tuple wrapping of the index argument

Example fix

# before
vals = ref.get((idx_a, idx_b), tree=tree)  # multiple indexers

# after
va = ref.get(idx_a)
vb = ref.get(idx_b)

indexers = jax.tree_util.tree_leaves(idx_tree)
assert len(indexers) <= 1, 'get supports at most one indexer'

def is_single_indexer(idx) -> bool:
    return len(jax.tree_util.tree_leaves(idx)) <= 1

When it happens

Trigger: Calling ref.get(...) (or indexing a Ref) where tree_unflatten of the indexer args yields len(indexers_avals) > 1 — i.e. multiple indexers passed to a single get_p bind, often from tuple/stacked index expressions inside a pallas kernel.

Common situations: Passing a tuple of indices or a pytree of indexers to a Ref get; using helper code that packs several indexers into one get call inside pallas kernels with block specs.

Related errors

• get not supported yet for block shape {b}
• Explicit sharding is not currently supported in the pallas-t
• group_offset is not currently supported in the pallas-triton
• Concatenation with Element indexing is not yet supported.
• Stack with Element indexing is not yet supported.

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