jax-ml/jax · error · ValueError

a_scale and b_scale must both be present or absent.

Error message

a_scale and b_scale must both be present or absent.

What it means

tcgen05.mma requires a_scale and b_scale to be passed together (block-scaled MMA) or omitted together. Passing only one raises this ValueError.

Source

Thrown at jax/_src/pallas/mosaic_gpu/primitives.py:2488

        acc.transforms)
    acc = acc.ref
  else:
    acc_transforms_leaves, acc_transforms_tree = [], None

  if isinstance(a, pallas_core.TransformedRef):
    a_transforms_leaves, a_transforms_tree = jax.tree.flatten(a.transforms)
    a = a.ref
  else:
    a_transforms_leaves, a_transforms_tree = [], None

  if isinstance(b, pallas_core.TransformedRef):
    b_transforms_leaves, b_transforms_tree = jax.tree.flatten(b.transforms)
    b = b.ref
  else:
    b_transforms_leaves, b_transforms_tree = [], None

  if (is_scaled := a_scale is not None) != (b_scale is not None):
    raise ValueError("a_scale and b_scale must both be present or absent.")
  scales = []
  if isinstance(a_scale, pallas_core.TransformedRef):
    a_scale_transforms_leaves, a_scale_transforms_tree = jax.tree.flatten(
        a_scale.transforms
    )
    scales.append(a_scale.ref)
  else:
    a_scale_transforms_leaves, a_scale_transforms_tree = [], None
    scales.append(a_scale)
  if isinstance(b_scale, pallas_core.TransformedRef):
    b_scale_transforms_leaves, b_scale_transforms_tree = jax.tree.flatten(
        b_scale.transforms
    )
    scales.append(b_scale.ref)
  else:
    b_scale_transforms_leaves, b_scale_transforms_tree = [], None
    scales.append(b_scale)
  if not is_scaled:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pass both a_scale and b_scale, or neither
2. If only one operand needs scaling conceptually, pass an all-ones scale ref for the other side
3. Gate both scale arguments on the same condition (e.g. 'if use_mxfp8: ... both')

Example fix

# before
tcgen05.mma(a, b, acc, k_dim=k, a_scale=a_s)
# after
tcgen05.mma(a, b, acc, k_dim=k, a_scale=a_s, b_scale=b_s)
# or omit both scales entirely

if (a_scale is None) != (b_scale is None):
    raise ValueError('need both scales or none')
tcgen05.mma(a, b, acc, k_dim=k, a_scale=a_scale, b_scale=b_scale)

When it happens

Trigger: Calling tcgen05.mma(a, b, acc, k_dim=k, a_scale=s) without b_scale, or vice versa, when experimenting with MXFP8/MXFP4 block scaling.

Common situations: Partially wiring up scale refs in a scaled-attention kernel; refactoring scale plumbing and dropping one argument; conditional code paths that supply scales asymmetrically.

Related errors

• a_scale must be a TMEM Ref
• b_scale must be a TMEM Ref
• shift must be non-negative.
• stride and stride_axis must be both specified or not.
• stride must be non-negative.

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