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.

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: "Must be a positive integer", "Invalid value", "Unsupported": the invalid-argument-value error family, when a library rejects the value you pass — this error's family across 35 libraries.

Related errors

• {ctx.avals_out[0].dtype}
• Expected an accumulator ref, got
• Loading from an accumulator is not supported. Use…
• Only float32 and int32 results are supported, got
• The shape of the accumulator

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

Appendix: 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)