{"record":{"id":"1c345149501a6bdc","repo":"jax-ml/jax","slug":"preinitialized-wgmmaaccumulatorref-only-supported","errorCode":null,"errorMessage":"Preinitialized WGMMAAccumulatorRef only supported in pl.run_state.","messagePattern":"Preinitialized WGMMAAccumulatorRef only supported in pl\\.run_state\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/core.py","lineNumber":1556,"sourceCode":"  def get_ref_aval(self) -> state.AbstractRef:\n    ty = ClusterBarrierType(\n        collective_axes=self.collective_axes,\n        num_arrivals=self.num_arrivals,\n        orders_tensor_core=self.orders_tensor_core,\n        leader_tracked=self.leader_tracked,\n    )\n    return state.AbstractRef(jax_core.ShapedArray(self.num_barriers, ty), SMEM)\n\n\n@dataclasses.dataclass(frozen=True)\nclass WGMMAAccumulatorRef:\n  shape: tuple[int, int]\n  dtype: jnp.dtype = jnp.float32\n  _init: Any = state_types.uninitialized\n\n  def get_ref_aval(self) -> state.AbstractRef:\n    if self._init is not state_types.uninitialized:\n      raise ValueError(\n          \"Preinitialized WGMMAAccumulatorRef only supported in pl.run_state.\"\n      )\n    return WGMMAAbstractAccumulatorRef(\n        jax_core.ShapedArray(shape=self.shape, dtype=self.dtype), MemorySpace.REGS\n    )\n\n  @staticmethod\n  def init(array):\n    return WGMMAAccumulatorRef(array.shape, array.dtype, _init=array)\n\n\ndef _wgmma_ref_type_mapping(ref: WGMMAAccumulatorRef):\n  aval = WGMMAAbstractAccumulatorRef(\n      jax_core.ShapedArray(shape=ref.shape, dtype=ref.dtype), MemorySpace.REGS\n  )\n  return aval, ref._init\nstate_types._ref_type_aval_mappings[WGMMAAccumulatorRef] = _wgmma_ref_type_mapping\n","sourceCodeStart":1538,"sourceCodeEnd":1574,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/core.py#L1538-L1574","documentation":"WGMMAAccumulatorRef only supports a preinitialized initial value (init=...) inside pl.run_state; anywhere else, get_ref_aval rejects a non-uninitialized _init because preinitialized accumulator refs aren't representable as plain register arrays.","triggerScenarios":"Constructing WGMMAAccumulatorRef(..., init=some_value) and then obtaining its ref aval outside pl.run_state — e.g. via a JAX transform tracing the kernel, or building the aval manually.","commonSituations":"Using init= to seed WGMMA accumulators in Mosaic kernels not wrapped in pl.run_state; tracing/transforming kernels that contain preinitialized accumulators.","solutions":["Wrap the kernel in pl.run_state so preinitialized accumulator refs are supported","Drop the init argument (leave the accumulator uninitialized) if run_state isn't used","Initialize the accumulator by a separate explicit copy/add inside the kernel"],"exampleFix":"// before\nacc_ref = plgpu.WGMMAAccumulatorRef((m, n), init=x)  # aval outside run_state -> ValueError\n// after\n@pl.run_state\ndef kernel(...):\n    acc_ref = plgpu.WGMMAAccumulatorRef((m, n), init=x)\n    ...","handlingStrategy":"validation","validationCode":"if init is not None:\n    assert in_run_state_context, 'WGMMA init= requires pl.run_state'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Wrap kernels using init= accumulators in pl.run_state","Omit init= outside run_state and initialize explicitly in-kernel"],"tags":["jax","pallas","wgmma","accumulator","api-misuse"],"backgroundTag":"unsupported-initialization-context","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}