jax-ml/jax · error · ValueError

Expected both or neither of scales to be specified.

Error message

Expected both or neither of scales to be specified.

What it means

tcgen05 MMA lowering requires that a_scale and b_scale are either both provided (scaled MMA, e.g. FP8 with MX formats) or both omitted. Supplying only one is ambiguous and rejected.

Source

Thrown at jax/experimental/mosaic/gpu/dialect_lowering.py:2515

) -> Sequence[ir.Value]:
  ctx.check_collective(op)

  def tmem_layout(operand):
    result = inference_utils.in_tmem_layout_for_operand(op, operand)
    # satisfy the type checker
    assert result is not None
    return result

  acc_ref = _tmem_ref_from_ir(op.accumulator, tmem_layout(op.accumulator))

  if op.a_sparse_metadata is not None:
    a_sparse_metadata = _tmem_ref_from_ir(
        op.a_sparse_metadata, tmem_layout(op.a_sparse_metadata)
    )
  else:
    a_sparse_metadata = None
  if (scaled := op.a_scale is not None) != (op.b_scale is not None):
    raise ValueError("Expected both or neither of scales to be specified.")
  if scaled:
    a_scale = _tmem_ref_from_ir(op.a_scale, tmem_layout(op.a_scale))  # pyrefly: ignore[bad-argument-type]
    b_scale = _tmem_ref_from_ir(op.b_scale, tmem_layout(op.b_scale))  # pyrefly: ignore[bad-argument-type]
  else:
    a_scale = None
    b_scale = None

  if utils.is_smem_ref(op.a):
    a_transforms, b_transforms = inference_utils.in_transforms(op)
    a_swizzle = swizzle_from_transforms_attr(a_transforms)
    b_swizzle = swizzle_from_transforms_attr(b_transforms)
    a_ref = unwrap_transformed_memref(op.a, a_transforms)
    b_ref = unwrap_transformed_memref(op.b, b_transforms)
  else:
    a_ref = _tmem_ref_from_ir(op.a, tmem_layout(op.a))
    [b_transforms] = inference_utils.in_transforms(op)
    b_swizzle = swizzle_from_transforms_attr(b_transforms)
    a_swizzle = b_swizzle

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pass both a_scale and b_scale, or neither
2. Audit wrapper functions for default None scale arguments leaking through
3. If only one side needs scaling, use scale=1 explicit constant scale tensor for the other side if the op semantics allow

Example fix

// before
mgpu.tcgen05_mma(a, b, acc, a_scale=sa)  # b_scale missing
// after
mgpu.tcgen05_mma(a, b, acc, a_scale=sa, b_scale=sb)

assert (a_scale is None) == (b_scale is None), 'provide both scales or neither'

def valid_scales(a_scale, b_scale) -> bool:
    return (a_scale is None) == (b_scale is None)

When it happens

Trigger: Calling tcgen05_mma with a_scale set but b_scale None, or vice versa — typically when building scaled matmuls with per-operand scale tensors.

Common situations: FP8/MXFP8 matmul kernels where the B operand's scale is accidentally dropped (e.g. None default from a wrapper) or where a user assumes scales are per-operand optional.

Related errors

• Sparse MMA not supported for M=64
• Swizzle={b_swizzle} is too big for MMA with M=64. Try loweri
• Only M=128 and M=64 are supported for MMA, but got M={m}
• Sparse MMA unsupported for f32
• MMA with element type {elem_type_str} does not support block

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