jax-ml/jax · error · ValueError

a and b must be 2D, but got: {a_type.shape} and {b_type.shap

Error message

a and b must be 2D, but got: {a_type.shape} and {b_type.shape}

What it means

This internal check in the Pallas Triton dot lowering requires both operands to be rank-2 (2D) tensors when it decomposes shapes into (m, k) and (k, n). Higher-rank batched dots are not handled here and raise ValueError with the offending shapes.

Source

Thrown at jax/_src/pallas/triton/lowering.py:2403

  elif isinstance(precision, tuple):
    a_precision, b_precision = precision
    if a_precision in _TF32_PRECISIONS or b_precision in _TF32_PRECISIONS:
      input_precision = tt_dialect.InputPrecision.TF32
    elif a_aval.dtype == jnp.float32:
      input_precision = tt_dialect.InputPrecision.IEEE
    else:
      input_precision = None

    acc_dtype = out_aval.dtype
    if acc_dtype not in (jnp.int32, jnp.float16, jnp.float64):
      acc_dtype = jnp.float32
  else:
    raise NotImplementedError(f"Unsupported dot precision: {precision}.")

  a_type = ir.RankedTensorType(a.type)
  b_type = ir.RankedTensorType(b.type)
  if len(a_type.shape) != 2 or len(b_type.shape) != 2:
    raise ValueError("a and b must be 2D, but got:"
                     f" {a_type.shape} and {b_type.shape}")

  m, k = a_type.shape
  _, n = b_type.shape
  if a_type.element_type == ir.F64Type.get():
    # Triton's MMAv2 fp64 path uses the m8n8k4 PTX instruction but aggregates
    # it with NumRegisters={m:2, n:1, k:4}, producing an effective m16n8k16
    # per-warp tile.  Blocks smaller than these minimums cause repM/repN/repK
    # to round to zero, corrupting the ValueTable and segfaulting the compiler.
    #   M >= 16  (2 × instrM=8)
    #   N >=  8  (1 × instrN=8)
    #   K >= 16  (4 × instrK=4)
    errors = []
    if m < 16:
      errors.append(f"M={m} < 16")
    if n < 8:
      errors.append(f"N={n} < 8")
    if k < 16:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Reshape operands to 2D before the dot: a.reshape(m, k) @ block-level, or vmap/loop over batch via the kernel grid
Use jnp.dot or lax.dot_general outside pallas for batched matmuls
Ensure block specs give each dot exactly 2D tiles

Example fix

# before
acc = pl.dot(a3d, b3d)  # shapes (B, M, K), (B, K, N)

# after
# launch grid with an extra dim and index 2D tiles
acc = pl.dot(a3d[i], b3d[i])

Defensive patterns

Strategy: validation

Validate before calling

def is_2d(a) -> bool:
    return getattr(a, 'ndim', 0) == 2
assert is_2d(a) and is_2d(b), f'{a.shape} {b.shape}'

Type guard

def is_2d_tile(a) -> bool:
    return a.ndim == 2

Prevention

Reshape inputs to 2D blocks in kernel setup; use grid dims instead of batched dots

When it happens

Trigger: Passing 3D/4D arrays (e.g. (batch, m, k)) directly to the pallas dot lowering without first reshaping/batching manually; using pl.dot inside a kernel whose operands are 1D vectors.

Common situations: Writing a batched matmul kernel by reusing lax.dot_general semantics instead of looping over the grid; forgetting that Mosaic kernels require explicit 2D block operands.

Related errors

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