jax-ml/jax · error · TypeError

load_staged_rhs must be an integer or None.

Error message

load_staged_rhs must be an integer or None.

What it means

matmul_acc_lhs accepts load_staged_rhs as an integer staging-register index or None. Passing True/False is explicitly rejected because True silently means 'register 1', which is almost never what the user intended.

Source

Thrown at jax/_src/pallas/mosaic/primitives.py:1272

  If `load_staged_rhs` is not None, the previously pushed RHS will be loaded
  from the given staging register _before_ the matrix multiplication begins.
  The results of the multiplication are accumulated into the specified
  accumulator slice.

  The MXU index is deduced from the provided accumulator.

  Args:
    acc: The accumulator slice used for results.
    lhs: The left-hand side operand. Must be M x 256. For M divisible by the
      number of sublanes multiplied by datatype packing.
    load_staged_rhs: The staging register to load the RHS from. If None, the RHS
      is not loaded from staging and the matmul will reuse the existing one.
  """
  # This is a common error. You might intend to say to load the staged RHS, but
  # True is equivalent to saying "load the staged RHS FROM REGISTER 1", which is
  # probably not what you intended.
  if isinstance(load_staged_rhs, bool):
    raise TypeError("load_staged_rhs must be an integer or None.")
  acc_ref, acc_transforms = sp.get_ref_and_transforms(acc, None, "matmul_acc_lhs")
  flat_acc_transforms, acc_transforms_treedef = tree_util.tree_flatten(
      acc_transforms
  )
  matmul_acc_lhs_p.bind(
      acc_ref,
      lhs,
      *flat_acc_transforms,
      acc_transforms_tree=acc_transforms_treedef,
      load_staged_rhs=load_staged_rhs,
  )


@matmul_acc_lhs_p.def_effectful_abstract_eval
def _matmul_acc_lhs_abstract_eval(
    acc: state.AbstractRef, lhs, *flat_acc_transforms, acc_transforms_tree, load_staged_rhs
):
  del load_staged_rhs,  # Unused.

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pass None if you do not want the staged RHS loaded
2. Pass the staging register index (e.g. 1 or 2) as an int if you do

Example fix

# before
matmul_acc_lhs(acc, lhs, load_staged_rhs=True)
# after
matmul_acc_lhs(acc, lhs, load_staged_rhs=1)  # or None

assert load_staged_rhs is None or (isinstance(load_staged_rhs, int) and not isinstance(load_staged_rhs, bool))

def valid_load_staged_rhs(v) -> bool:
    return v is None or (isinstance(v, int) and not isinstance(v, bool))

When it happens

Trigger: Calling matmul_acc_lhs(acc, lhs, load_staged_rhs=True) or with False instead of an integer index or None.

Common situations: Intending 'yes, load the staged RHS' and writing True; new users assuming a boolean flag API from other pallas_matmul helpers.

Understand the failure class

Background: Invalid argument type errors: "must be of type string", "expected X, got Y", and ERR_INVALID_ARG_TYPE explained — this error's family across 15 libraries.

Related errors

• Loading from an accumulator is not supported. Use `matmul_po
• {ctx.avals_out[0].dtype}
• Transforms not supported for matmul_acc_lhs.
• Transforms not supported for matmul_pop.
• Expected an accumulator ref, got {acc}

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