{"record":{"id":"e16a701d69413531","repo":"jax-ml/jax","slug":"load-staged-rhs-must-be-an-integer-or-none","errorCode":null,"errorMessage":"load_staged_rhs must be an integer or None.","messagePattern":"load_staged_rhs must be an integer or None\\.","errorType":"exception","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/primitives.py","lineNumber":1272,"sourceCode":"  If `load_staged_rhs` is not None, the previously pushed RHS will be loaded\n  from the given staging register _before_ the matrix multiplication begins.\n  The results of the multiplication are accumulated into the specified\n  accumulator slice.\n\n  The MXU index is deduced from the provided accumulator.\n\n  Args:\n    acc: The accumulator slice used for results.\n    lhs: The left-hand side operand. Must be M x 256. For M divisible by the\n      number of sublanes multiplied by datatype packing.\n    load_staged_rhs: The staging register to load the RHS from. If None, the RHS\n      is not loaded from staging and the matmul will reuse the existing one.\n  \"\"\"\n  # This is a common error. You might intend to say to load the staged RHS, but\n  # True is equivalent to saying \"load the staged RHS FROM REGISTER 1\", which is\n  # probably not what you intended.\n  if isinstance(load_staged_rhs, bool):\n    raise TypeError(\"load_staged_rhs must be an integer or None.\")\n  acc_ref, acc_transforms = sp.get_ref_and_transforms(acc, None, \"matmul_acc_lhs\")\n  flat_acc_transforms, acc_transforms_treedef = tree_util.tree_flatten(\n      acc_transforms\n  )\n  matmul_acc_lhs_p.bind(\n      acc_ref,\n      lhs,\n      *flat_acc_transforms,\n      acc_transforms_tree=acc_transforms_treedef,\n      load_staged_rhs=load_staged_rhs,\n  )\n\n\n@matmul_acc_lhs_p.def_effectful_abstract_eval\ndef _matmul_acc_lhs_abstract_eval(\n    acc: state.AbstractRef, lhs, *flat_acc_transforms, acc_transforms_tree, load_staged_rhs\n):\n  del load_staged_rhs,  # Unused.","sourceCodeStart":1254,"sourceCodeEnd":1290,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/primitives.py#L1254-L1290","documentation":"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.","triggerScenarios":"Calling matmul_acc_lhs(acc, lhs, load_staged_rhs=True) or with False instead of an integer index or None.","commonSituations":"Intending 'yes, load the staged RHS' and writing True; new users assuming a boolean flag API from other pallas_matmul helpers.","solutions":["Pass None if you do not want the staged RHS loaded","Pass the staging register index (e.g. 1 or 2) as an int if you do"],"exampleFix":"# before\nmatmul_acc_lhs(acc, lhs, load_staged_rhs=True)\n# after\nmatmul_acc_lhs(acc, lhs, load_staged_rhs=1)  # or None","handlingStrategy":"type-guard","validationCode":"assert load_staged_rhs is None or (isinstance(load_staged_rhs, int) and not isinstance(load_staged_rhs, bool))","typeGuard":"def valid_load_staged_rhs(v) -> bool:\n    return v is None or (isinstance(v, int) and not isinstance(v, bool))","tryCatchPattern":null,"preventionTips":["Never pass booleans to load_staged_rhs; True means register 1"],"tags":["jax","pallas","tpu","matmul","type-error"],"backgroundTag":"invalid-argument-type","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}