jax-ml/jax · error · NotImplementedError

Uninitialized Refs are not supported in lowering of run_stat

Error message

Uninitialized Refs are not supported in lowering of run_state.

What it means

Identical condition to the run_state resource check: the run_state lowering requires all refs initialized (is_initialized all True). Uninitialized refs — refs with no bound contents at lowering time — are rejected with NotImplementedError.

Source

Thrown at jax/_src/pallas/mosaic_gpu/lowering.py:3746

    return ctx.module_ctx.reserve_semaphores(
        what.shape, collective_axes=collective_axes
    ).__enter__()
  raise NotImplementedError(f"get_global only supports semaphores, got {what}")


@register_lowering_rule(discharge.run_state_p, mgpu.LoweringSemantics.Lane)
@register_lowering_rule(discharge.run_state_p, mgpu.LoweringSemantics.Warpgroup)
def _run_state_lowering_rule(
    ctx: LoweringRuleContext,
    *args,
    jaxpr: jax_core.Jaxpr,
    which_linear: tuple[bool, ...],
    is_initialized: tuple[bool, ...],
):
  del which_linear
  # TODO(apaszke): This should be unified with run_scoped.
  if not all(is_initialized):
    raise NotImplementedError("Uninitialized Refs are not supported in lowering of run_state.")

  should_discharge = []
  new_input_vals = []
  # `should_deref_acc` is used under lane lowering semantics, to figure out
  # whether we need to return a `WGMMAAccumulator` or a `FragmentedArray` when
  # encountering a `WGMMAAbstractAccumulatorRef` as input.
  #
  # We can't tell the difference under warpgroup lowering semantics, but we do
  # not need to since we always return a `vector` anyway.
  should_deref_acc = []
  for arg, v, out_aval in zip(args, jaxpr.invars, ctx.avals_out):
    aval = v.aval
    if isinstance(aval, gpu_core.WGMMAAbstractAccumulatorRef):
      if ctx.module_ctx.lowering_semantics == mgpu.LoweringSemantics.Warpgroup:
        arg = mgpu.dialect.optimization_barrier([arg])
        nvvm_dialect.wgmma_fence_aligned()
        new_input_vals.append(arg)
      else:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Initialize every ref (store a zero/default) at scope entry before conditional logic
2. Hoist conditional stores so each ref is written exactly once unconditionally
3. Reproduce on the latest jax version and file an issue if all refs are provably initialized

Example fix

// before
pl.run_scoped(lambda r: maybe_store(r), ref)
// after
pl.run_scoped(lambda r: (r[...].set(0), maybe_store(r))[1], ref)

for ref in refs:
    ref[...] = zeros  # unconditional init before any conditional stores

When it happens

Trigger: discharge.run_state lowering encountering a ref flagged uninitialized, e.g. refs created but never stored to before being discharged out of a scope.

Common situations: Scoped allocations whose body may skip stores on some paths; partial initialization under data-dependent control flow in kernels; internal discharge changes across JAX versions.

Related errors

• This gmm kernel only supports either (m, k) x (g, k, n) -> (
• Group sizes {group_sizes.shape=} must match first dimension
• Explicit sharding is not currently supported in the pallas-t
• group_offset is not currently supported in the pallas-triton
• Compiler params for platform {platform} cannot be used for {

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