{"record":{"id":"51f26ec901b4e271","repo":"jax-ml/jax","slug":"existing-output-dtype-must-match-preferred-element","errorCode":null,"errorMessage":"Existing output dtype must match preferred_element_type.","messagePattern":"Existing output dtype must match preferred_element_type\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/pallas/ops/tpu/megablox/gmm.py","lineNumber":348,"sourceCode":"    group_sizes: A 1d, jnp.ndarray with shape [num_groups] and jnp.int32 dtype.\n    preferred_element_type: jnp.dtype, the element type for the output matrix.\n    tiling: 3-tuple of ints. The m, k and n-dimension tile sizes.\n    group_offset: The group in group sizes to start computing from. This is\n      particularly useful for when rhs num_groups is sharded.\n    existing_out: Existing output to write to.\n    transpose_rhs: True if the rhs needs to be transposed.\n    interpret: Whether or not to run the kernel in interpret mode, helpful for\n      testing and debugging.\n\n  Returns:\n    A 2d, jnp.ndarray with shape [m, n].\n  \"\"\"\n\n  if existing_out is not None:\n    assert isinstance(existing_out, jax.Array)\n    expected_dtype = existing_out.dtype\n    if expected_dtype != preferred_element_type:\n      raise ValueError(\n          \"Existing output dtype must match preferred_element_type.\"\n      )\n  if group_offset is None:\n    group_offset = jnp.array([0], dtype=jnp.int32)\n  else:\n    if group_offset.shape:\n      raise ValueError(\n          f\"group_offset must be a ()-shaped array. Got: {group_offset.shape}.\"\n      )\n    group_offset = group_offset[None]\n  num_current_groups = rhs.shape[0]\n  num_total_groups = group_sizes.shape[0]\n  lhs, group_sizes, input_dtype = _validate_args(\n      lhs=lhs, rhs=rhs, group_sizes=group_sizes\n  )\n\n  # Gather shape information.\n  m, k, n = (lhs.shape[0], lhs.shape[1], rhs.shape[2])","sourceCodeStart":330,"sourceCodeEnd":366,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/pallas/ops/tpu/megablox/gmm.py#L330-L366","documentation":"Thrown by jax.experimental.pallas.ops.tpu.megablox.gmm when an existing_out array is supplied whose dtype differs from the preferred_element_type argument. Megablox's GMM reuses an existing output buffer for accumulation, and the kernel requires the buffer dtype to exactly match the requested output element type. Any mismatch aborts with this ValueError before kernel launch.","triggerScenarios":"Calling gmm(lhs, rhs, group_sizes, existing_out=out_buf, preferred_element_type=dt) where out_buf.dtype != dt, e.g. existing_out float32 buffer with preferred_element_type=jnp.bfloat16. Also triggered indirectly via _gmm_fwd/_gmm_bwd autodiff paths that pass existing_out.","commonSituations":"Mixed-precision training where activations are bf16 but an fp32 output buffer is reused; refactoring code that previously did not pass preferred_element_type; upgrading JAX versions where megablox began enforcing the check.","solutions":["Set preferred_element_type to exactly existing_out.dtype (or vice versa), e.g. preferred_element_type=existing_out.dtype","Cast the reused buffer: existing_out = existing_out.astype(preferred_element_type)","Drop existing_out entirely if you do not need in-place accumulation"],"exampleFix":"// before\nout = gmm(lhs, rhs, group_sizes, existing_out=buf, preferred_element_type=jnp.bfloat16)\n// after\nout = gmm(lhs, rhs, group_sizes, existing_out=buf, preferred_element_type=buf.dtype)","handlingStrategy":"validation","validationCode":"assert existing_out is None or existing_out.dtype == preferred_element_type, (existing_out.dtype, preferred_element_type)","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Derive preferred_element_type from the buffer: preferred_element_type=existing_out.dtype","Add a unit test asserting dtype agreement for mixed-precision paths"],"tags":["jax","pallas","tpu","megablox","dtype-mismatch"],"backgroundTag":"dtype-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}