jax-ml/jax · error · NotImplementedError

Grid mapping with hijax index maps are not currently support

Error message

Grid mapping with hijax index maps are not currently supported. Got {grid_mapping=}

What it means

_emit_pipeline_is_high checks whether the pipeline primitive must run in high (hijax) mode. Grid mappings whose block index_map jaxprs (or their constants) are themselves high-level hijax values are not supported, so NotImplementedError is raised with the offending grid_mapping.

Source

Thrown at jax/_src/pallas/mosaic/pipeline.py:2626

  lowering_context = ctx.lowering_context.replace(
      block_shapes=ctx.block_shapes,
      grid_sizes=grid_sizes,
      grid_names=grid_names,
      user_grid_indices=user_grid_indices,
      vmapped_dims=grid_mapping.vmapped_dims,
      emit_pipeline_mode=True,
  )

  assert len(jaxpr.invars) == len(lowering_context.block_shapes)
  assert len(lowering_context.grid_sizes) == len(lowering_context.grid_names)
  return jaxpr_subcomp(lowering_context, jaxpr, *args_flat)

def _emit_pipeline_is_high(*avals, body_jaxpr, grid_mapping, args_tree, **_):
  # Check that the index_maps jaxpr or consts are not high.
  if (any(bm.index_map_jaxpr.is_high for bm in grid_mapping.block_mappings)
      or any(any(c.is_high for c in bm.index_map_jaxpr.consts)
             for bm in grid_mapping.block_mappings)):
    raise NotImplementedError("Grid mapping with hijax index maps are not"
                              f" currently supported. Got {grid_mapping=}")

  return (body_jaxpr.is_high
          or any(bm.transformed_block_aval.inner_aval.is_high
                 for bm in grid_mapping.block_mappings))

emit_pipeline_p.is_high = _emit_pipeline_is_high


def _emit_pipeline_to_lojax(
    *args_flat, body_jaxpr, grid_mapping, args_tree, refs_tree, **params
):
  all_args: EmitPipelinePrimitiveArgs = args_tree.unflatten(args_flat)
  closed_hi_jaxpr = core.ClosedJaxpr(body_jaxpr, all_args.body_consts)
  with grid_mapping.trace_env():
    closed_lo_jaxpr = pe.lower_jaxpr2(closed_hi_jaxpr)

  refs_avals = [jax.typeof(x) for x in all_args.refs_flat]

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Keep index maps as plain low-level callables over ordinary JAX values (no hijax types captured)
2. Rebuild the BlockSpec so its index_map traces to a low jaxpr
3. If you believe it should work, report to the jax repo with the grid_mapping printout

try:
    emit_pipeline(...)
except NotImplementedError as e:
    if 'hijax index maps' in str(e):
        # rebuild BlockSpec index maps as plain low-level callables
        ...

When it happens

Trigger: Constructing a pipeline whose BlockSpec index_map closes over or produces hijax (high-level) values — e.g. using hijax features/tracers inside index maps rather than low-level closed jaxprs.

Common situations: Mixing experimental hijax frontend objects into index maps; recent JAX version changes in how index maps are traced; advanced users composing index maps programmatically.

Understand the failure class

Background: UnsupportedOperationException and "is not supported" errors: when a library deliberately refuses a call — this error's family across 30 libraries.

Related errors

• Explicit sharding is not currently supported in the pallas-t
• group_offset is not currently supported in the pallas-triton
• Index map function {debug_info.func_src_info} for {origin} m
• index_map returned a value of type {type(idx_aval)} at posit
• index_map returned a value of type {type(idx_aval)} at posit

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