{"record":{"id":"4db21d05c4447c1f","repo":"jax-ml/jax","slug":"uninitialized-refs-are-not-supported-in-lowering-o","errorCode":null,"errorMessage":"Uninitialized Refs are not supported in lowering of run_state.","messagePattern":"Uninitialized Refs are not supported in lowering of run_state\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/lowering.py","lineNumber":3746,"sourceCode":"    return ctx.module_ctx.reserve_semaphores(\n        what.shape, collective_axes=collective_axes\n    ).__enter__()\n  raise NotImplementedError(f\"get_global only supports semaphores, got {what}\")\n\n\n@register_lowering_rule(discharge.run_state_p, mgpu.LoweringSemantics.Lane)\n@register_lowering_rule(discharge.run_state_p, mgpu.LoweringSemantics.Warpgroup)\ndef _run_state_lowering_rule(\n    ctx: LoweringRuleContext,\n    *args,\n    jaxpr: jax_core.Jaxpr,\n    which_linear: tuple[bool, ...],\n    is_initialized: tuple[bool, ...],\n):\n  del which_linear\n  # TODO(apaszke): This should be unified with run_scoped.\n  if not all(is_initialized):\n    raise NotImplementedError(\"Uninitialized Refs are not supported in lowering of run_state.\")\n\n  should_discharge = []\n  new_input_vals = []\n  # `should_deref_acc` is used under lane lowering semantics, to figure out\n  # whether we need to return a `WGMMAAccumulator` or a `FragmentedArray` when\n  # encountering a `WGMMAAbstractAccumulatorRef` as input.\n  #\n  # We can't tell the difference under warpgroup lowering semantics, but we do\n  # not need to since we always return a `vector` anyway.\n  should_deref_acc = []\n  for arg, v, out_aval in zip(args, jaxpr.invars, ctx.avals_out):\n    aval = v.aval\n    if isinstance(aval, gpu_core.WGMMAAbstractAccumulatorRef):\n      if ctx.module_ctx.lowering_semantics == mgpu.LoweringSemantics.Warpgroup:\n        arg = mgpu.dialect.optimization_barrier([arg])\n        nvvm_dialect.wgmma_fence_aligned()\n        new_input_vals.append(arg)\n      else:","sourceCodeStart":3728,"sourceCodeEnd":3764,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/lowering.py#L3728-L3764","documentation":"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.","triggerScenarios":"discharge.run_state lowering encountering a ref flagged uninitialized, e.g. refs created but never stored to before being discharged out of a scope.","commonSituations":"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.","solutions":["Initialize every ref (store a zero/default) at scope entry before conditional logic","Hoist conditional stores so each ref is written exactly once unconditionally","Reproduce on the latest jax version and file an issue if all refs are provably initialized"],"exampleFix":"// before\npl.run_scoped(lambda r: maybe_store(r), ref)\n// after\npl.run_scoped(lambda r: (r[...].set(0), maybe_store(r))[1], ref)","handlingStrategy":"validation","validationCode":"for ref in refs:\n    ref[...] = zeros  # unconditional init before any conditional stores","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Write defaults to refs immediately after allocation","Make every ref store unconditional or preceded by an init","Report suspected discharge bugs with a minimal reproducer"],"tags":["jax","pallas","run-state","uninitialized-ref"],"backgroundTag":"uninitialized-value","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}