jax-ml/jax · error · NotImplementedError

get not supported yet for block shape {b}

Error message

get not supported yet for block shape {b}

What it means

Raised by the _slice helper inside _get_eval_rule when the ref's block_shape entry for a dimension is not int, pallas_core.Blocked, pallas_core.Squeezed, or None — typically pallas_core.Element. The eval rule cannot compute the per-block index (i * b) for such block descriptors when materializing a get.

Source

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

    raise NotImplementedError('get not supported yet')
  if not indexers:
    indexer = indexing.NDIndexer.make_trivial_indexer(ref_aval.shape)
    indexer_aval = indexer
  else:
    indexer = indexers[0]
    indexer_aval = indexers_avals[0]
  block_indexer = []

  def _slice(i, b):
    match b:
      case int():
        return indexing.ds(i * b, b)
      case pallas_core.Blocked(bs):
        return indexing.ds(i * bs, bs)
      case pallas_core.Squeezed() | None:
        return i
      case _:
        raise NotImplementedError(f'get not supported yet for block shape {b}')

  if (
      ref_block_spec is pallas_core.no_block_spec
      or ref_block_spec.block_shape is None
  ):
    # Short-circuit if the ref is not blocked.
    return state_primitives.get_p.bind(ref, *idx, tree=tree)
  block_idx_iter = iter(ctx.get_out_block_indices()[0])
  for idx_aval, size, idx, bd in zip(
      indexer_aval.indices,
      ref_aval.shape,
      indexer.indices,
      ref_block_spec.block_shape,
      strict=True,
  ):
    if not isinstance(idx_aval, indexing.Slice):
      assert hasattr(idx_aval, 'shape') and not idx_aval.shape, idx_aval
      assert bd is None or isinstance(bd, pallas_core.Squeezed)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Change the BlockSpec to use int/Blocked/Squeezed block descriptors on dims that get read
2. Avoid get on Element-blocked refs; read via an int-blocked alias buffer

Example fix

# before
BlockSpec(block_shape=(pallas_core.Element(1), 64), ...)
v = ref[:]

# after
BlockSpec(block_shape=(1, 64), ...)
v = ref[:]

import jax._src.pallas.pallas_core as pc
assert all(b is None or isinstance(b, (int, pc.Blocked, pc.Squeezed)) for b in block_shape), f'unsupported block type for get: {block_shape}'

def is_gettable_block(b) -> bool:
    import jax._src.pallas.pallas_core as pc
    return b is None or isinstance(b, (int, pc.Blocked, pc.Squeezed))

When it happens

Trigger: ref.get / ref[...] on a Ref blocked with pallas_core.Element (or any exotic block descriptor) inside a pallas kernel evaluated by the fuser.

Common situations: Using Element block specs for fine-grained layouts then reading the ref; mixing Element-blocked outputs with gets in fused kernels.

Related errors

• get not supported yet
• Concatenation with Element indexing is not yet supported.
• Stack with Element indexing is not yet supported.
• Explicit sharding is not currently supported in the pallas-t
• group_offset is not currently supported in the pallas-triton

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