{"record":{"id":"728a68ca1e090467","repo":"jax-ml/jax","slug":"grid-mapping-with-hijax-index-maps-are-not-current","errorCode":null,"errorMessage":"Grid mapping with hijax index maps are not currently supported. Got {grid_mapping=}","messagePattern":"Grid mapping with hijax index maps are not currently supported\\. Got (.+?)","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/pipeline.py","lineNumber":2626,"sourceCode":"  lowering_context = ctx.lowering_context.replace(\n      block_shapes=ctx.block_shapes,\n      grid_sizes=grid_sizes,\n      grid_names=grid_names,\n      user_grid_indices=user_grid_indices,\n      vmapped_dims=grid_mapping.vmapped_dims,\n      emit_pipeline_mode=True,\n  )\n\n  assert len(jaxpr.invars) == len(lowering_context.block_shapes)\n  assert len(lowering_context.grid_sizes) == len(lowering_context.grid_names)\n  return jaxpr_subcomp(lowering_context, jaxpr, *args_flat)\n\ndef _emit_pipeline_is_high(*avals, body_jaxpr, grid_mapping, args_tree, **_):\n  # Check that the index_maps jaxpr or consts are not high.\n  if (any(bm.index_map_jaxpr.is_high for bm in grid_mapping.block_mappings)\n      or any(any(c.is_high for c in bm.index_map_jaxpr.consts)\n             for bm in grid_mapping.block_mappings)):\n    raise NotImplementedError(\"Grid mapping with hijax index maps are not\"\n                              f\" currently supported. Got {grid_mapping=}\")\n\n  return (body_jaxpr.is_high\n          or any(bm.transformed_block_aval.inner_aval.is_high\n                 for bm in grid_mapping.block_mappings))\n\nemit_pipeline_p.is_high = _emit_pipeline_is_high\n\n\ndef _emit_pipeline_to_lojax(\n    *args_flat, body_jaxpr, grid_mapping, args_tree, refs_tree, **params\n):\n  all_args: EmitPipelinePrimitiveArgs = args_tree.unflatten(args_flat)\n  closed_hi_jaxpr = core.ClosedJaxpr(body_jaxpr, all_args.body_consts)\n  with grid_mapping.trace_env():\n    closed_lo_jaxpr = pe.lower_jaxpr2(closed_hi_jaxpr)\n\n  refs_avals = [jax.typeof(x) for x in all_args.refs_flat]","sourceCodeStart":2608,"sourceCodeEnd":2644,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/pipeline.py#L2608-L2644","documentation":"_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.","triggerScenarios":"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.","commonSituations":"Mixing experimental hijax frontend objects into index maps; recent JAX version changes in how index maps are traced; advanced users composing index maps programmatically.","solutions":["Keep index maps as plain low-level callables over ordinary JAX values (no hijax types captured)","Rebuild the BlockSpec so its index_map traces to a low jaxpr","If you believe it should work, report to the jax repo with the grid_mapping printout"],"exampleFix":null,"handlingStrategy":"fallback","validationCode":null,"typeGuard":null,"tryCatchPattern":"try:\n    emit_pipeline(...)\nexcept NotImplementedError as e:\n    if 'hijax index maps' in str(e):\n        # rebuild BlockSpec index maps as plain low-level callables\n        ...","preventionTips":["Keep index maps free of hijax/experimental high-level values","Pin a known-good JAX version when using advanced pipeline features"],"tags":["jax","pallas","hijax","index-map","not-implemented"],"backgroundTag":"unsupported-operation","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}