jax-ml/jax · error · ValueError

Existing output dtype must match preferred_element_type.

Error message

Existing output dtype must match preferred_element_type.

What it means

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.

Source

Thrown at jax/experimental/pallas/ops/tpu/megablox/gmm.py:348

    group_sizes: A 1d, jnp.ndarray with shape [num_groups] and jnp.int32 dtype.
    preferred_element_type: jnp.dtype, the element type for the output matrix.
    tiling: 3-tuple of ints. The m, k and n-dimension tile sizes.
    group_offset: The group in group sizes to start computing from. This is
      particularly useful for when rhs num_groups is sharded.
    existing_out: Existing output to write to.
    transpose_rhs: True if the rhs needs to be transposed.
    interpret: Whether or not to run the kernel in interpret mode, helpful for
      testing and debugging.

  Returns:
    A 2d, jnp.ndarray with shape [m, n].
  """

  if existing_out is not None:
    assert isinstance(existing_out, jax.Array)
    expected_dtype = existing_out.dtype
    if expected_dtype != preferred_element_type:
      raise ValueError(
          "Existing output dtype must match preferred_element_type."
      )
  if group_offset is None:
    group_offset = jnp.array([0], dtype=jnp.int32)
  else:
    if group_offset.shape:
      raise ValueError(
          f"group_offset must be a ()-shaped array. Got: {group_offset.shape}."
      )
    group_offset = group_offset[None]
  num_current_groups = rhs.shape[0]
  num_total_groups = group_sizes.shape[0]
  lhs, group_sizes, input_dtype = _validate_args(
      lhs=lhs, rhs=rhs, group_sizes=group_sizes
  )

  # Gather shape information.
  m, k, n = (lhs.shape[0], lhs.shape[1], rhs.shape[2])

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Set preferred_element_type to exactly existing_out.dtype (or vice versa), e.g. preferred_element_type=existing_out.dtype
2. Cast the reused buffer: existing_out = existing_out.astype(preferred_element_type)
3. Drop existing_out entirely if you do not need in-place accumulation

Example fix

// before
out = gmm(lhs, rhs, group_sizes, existing_out=buf, preferred_element_type=jnp.bfloat16)
// after
out = gmm(lhs, rhs, group_sizes, existing_out=buf, preferred_element_type=buf.dtype)

assert existing_out is None or existing_out.dtype == preferred_element_type, (existing_out.dtype, preferred_element_type)

When it happens

Trigger: 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.

Common situations: 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.

Related errors

• group_offset must be a ()-shaped array. Got: {group_offset.s
• No tuned tiling found for (m, k, n) = ({m}, {k}, {n})
• Compiler params for platform {platform} cannot be used for {
• Memory space {self.memory_space} is not supported by mesh {s
• Acc ref must be at least 2D, got shape {shape}

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