jax-ml/jax · error · ValueError

scaled_matmul requires all inputs to be 3-dimensional arrays

Error message

scaled_matmul requires all inputs to be 3-dimensional arrays

What it means

jax.nn.scaled_matmul implements block-scaled (MX-format float8) matrix multiplication and requires lhs, rhs, lhs_scales, rhs_scales all to be exactly 3D (B, M, K) / (B, N, K) blocks. Passing any 2D or 4D tensor raises this.

Source

Thrown at jax/_src/nn/functions.py:1355

      >>> a = jnp.array([1, 2, 3]).reshape((1, 1, 3))
      >>> b = jnp.array([4, 5, 6]).reshape((1, 1, 3))
      >>> a_scales = jnp.array([0.5]).reshape((1, 1, 1))
      >>> b_scales = jnp.array([0.5]).reshape((1, 1, 1))
      >>> scaled_matmul(a, b, a_scales, b_scales)  # doctest: +SKIP
      Array([[[8.]]], dtype=float32)

      Using fused cuDNN call on Blackwell GPUs:

      >>> dtype = jnp.float8_e4m3fn
      >>> a = jax.random.normal(jax.random.PRNGKey(1), (3, 128, 64), dtype=dtype)
      >>> b = jax.random.normal(jax.random.PRNGKey(2), (3, 128, 64), dtype=dtype)
      >>> a_scales = jnp.ones((3, 128, 4), dtype=jnp.float8_e8m0fnu)
      >>> b_scales = jnp.ones((3, 128, 4), dtype=jnp.float8_e8m0fnu)
      >>> scaled_matmul(a, b, a_scales, b_scales)  # doctest: +SKIP
    """
    a, b, a_scales, b_scales = lhs, rhs, lhs_scales, rhs_scales
    if not all(x.ndim == 3 for x in (a, b, a_scales, b_scales)):
        raise ValueError(
            "scaled_matmul requires all inputs to be 3-dimensional arrays"
        )

    B_a, M_a, K_a = a.shape
    B_b, N_b, K_b = b.shape
    if K_a != K_b or B_a != B_b:
        raise ValueError(
            "scaled_matmul requires inputs a and b to have matching batch (B) "
            f"and contract (K) dimensions, but got shapes {a.shape} and "
            f"{b.shape}"
        )

    B_as, M_as, K_as = a_scales.shape
    B_bs, N_bs, K_bs = b_scales.shape
    if K_as != K_bs or B_as != B_bs:
        raise ValueError(
            "scaled_matmul requires scales to have matching batch (B) and "
            f"contract (K) dimensions, but got shapes {a_scales.shape} and "

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Add a leading batch dimension: a[None], b[None], a_scales[None], b_scales[None]
2. Reshape 4D tensors to 3D (fold batch*heads) if appropriate
3. Build scales with the block layout (B, ceil(M/32), ceil(K/32)) as float8_e8m0fnu

Example fix

// before
out = jax.nn.scaled_matmul(a2d, b2d, a_s2d, b_s2d)

// after
out = jax.nn.scaled_matmul(a2d[None], b2d[None], a_s2d[None], b_s2d[None])[0]

assert all(x.ndim == 3 for x in (a, b, a_s, b_s)), 'scaled_matmul needs 3D (B,M,K)/(B,N,K) inputs'
if a.ndim == 2: a, b, a_s, b_s = (t[None] for t in (a, b, a_s, b_s))

def is_3d(*ts) -> bool: return all(getattr(t, 'ndim', -1) == 3 for t in ts)

When it happens

Trigger: Passing plain 2D matrices without an outer batch dim, or 4D attention tensors, to scaled_matmul; forgetting the scales arrays or passing scalars.

Common situations: Adapting a normal jnp.matmul call to MXFP8 scaled matmul and forgetting to add a batch dimension; passing per-tensor scalar scales instead of 3D block scales.

Related errors

• scaled_matmul requires inputs a and b to have matching batch
• scaled_matmul requires scales to have matching batch (B) and
• scaled_matmul requires scales to match non-contract dimensio
• {name} ndim should be {len(shape)}, but got {t.ndim}
• {name} shape should be {shape}: but got {t.shape}

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