{"record":{"id":"426322a4d8aee17c","repo":"jax-ml/jax","slug":"either-none-or-both-scales-should-be-provided","errorCode":null,"errorMessage":"Either none or both scales should be provided","messagePattern":"Either none or both scales should be provided","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/tcgen05.py","lineNumber":200,"sourceCode":"    b: ir.Value,\n    *,\n    a_swizzle: int = 128,\n    b_swizzle: int = 128,\n    a_scale: TMEMRef | None = None,\n    b_scale: TMEMRef | None = None,\n    a_sparse_metadata: TMEMRef | None = None,\n    accumulate: ir.Value | bool = True,\n    collective: bool = False,\n) -> None:\n  if a_swizzle == 16 or b_swizzle == 16:\n    raise NotImplementedError(\"No swizzle is not supported\")\n  i8 = ir.IntegerType.get_signless(8)\n  i32 = ir.IntegerType.get_signless(32)\n  if isinstance(accumulate, bool):\n    accumulate = arith.constant(ir.IntegerType.get_signless(1), accumulate)\n  num_cta = 2 if collective else 1\n  if (is_scaled := a_scale is not None) != (b_scale is not None):\n    raise ValueError(\"Either none or both scales should be provided\")\n  is_sparse = a_sparse_metadata is not None\n  if is_scaled and is_sparse:\n    if isinstance(a, TMEMRef):\n      raise NotImplementedError(\n          \"A in TMEM unsupported for block-scaled sparse matmuls\"\n      )\n\n  # Step 1. Establish the shape and element type of the operation.\n  if not isinstance(b.type, ir.MemRefType):\n    raise ValueError(f\"B must be a memref, got: {b.type}\")\n  (k, n), b_element_type = mma_utils.tiled_memref_shape(b)\n  if isinstance(a, TMEMRef):\n    m, k2 = a.shape\n    a_element_type = a.dtype\n    if m != 128:\n      raise NotImplementedError(f\"Only M=128 is supported for MMA with A in TMEM, but got M={m}\")\n    # Watch out: this layout must be consistent with D's layout (up to packing).\n    expected_packing = 32 // utils.bitwidth(a_element_type)","sourceCodeStart":182,"sourceCodeEnd":218,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/tcgen05.py#L182-L218","documentation":"Block-scaled MMA requires scale factors for both operands or neither — the hardware descriptor has symmetric A/B scale fields. Passing only a_scale (or only b_scale) is an API misuse and raises ValueError.","triggerScenarios":"Calling mma(..., a_scale=tref) without b_scale, or vice versa; conditionally passing one scale based on a None check.","commonSituations":"Refactoring code that had scales on only one operand in a different API; optional-argument plumbing where one scale fails to propagate.","solutions":["Pass both a_scale and b_scale, or omit both","If only one operand needs scaling, pass an identity/1.0-filled scale for the other operand"],"exampleFix":"# before\nmma(acc, a, b, a_scale=a_scale)\n\n# after\nmma(acc, a, b, a_scale=a_scale, b_scale=b_scale)  # or drop both","handlingStrategy":"validation","validationCode":"if (a_scale is None) != (b_scale is None):\n    raise ValueError('provide both a_scale and b_scale, or neither')\nmma(acc, a, b, a_scale=a_scale, b_scale=b_scale)","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Pass scales as a pair","Use a one-filled b_scale when only one operand is truly scaled"],"tags":["tcgen05","matmul","mx-format","argument-validation"],"backgroundTag":"mismatched-optional-arguments","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}