{"record":{"id":"f37c256236a7aafd","repo":"jax-ml/jax","slug":"a-scale-layout-a-scale-layout-is-not-supported","errorCode":null,"errorMessage":"A scale layout {a_scale.layout} is not supported","messagePattern":"A scale layout (.+?) is not supported","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/tcgen05.py","lineNumber":472,"sourceCode":"              f\" {scale_block}\"\n          )\n      else:\n        raise ValueError(\n            \"Scale element type mismatch: expected f8e8m0fnu or f8e4m3fn, got\"\n            f\" {scale_element_type}\"\n        )\n    else:\n      raise NotImplementedError(\n          f\"Unsupported element type for block scaling: {a_element_type}\"\n      )\n    k_scales = k // scale_block\n    if a_scale.shape != (TMEM_ROWS, k_scales):\n      raise ValueError(\n          f\"A scale shape mismatch: expected ({TMEM_ROWS}, {k_scales}), got\"\n          f\" {a_scale.shape}\"\n      )\n    if a_scale.layout != scales_layout():\n      raise ValueError(f\"A scale layout {a_scale.layout} is not supported\")\n    if collective and m == 64:\n      if b_scale.layout != b_scales_m64_collective_layout():\n        raise ValueError(\n            \"Expected B scales to have a M=64 collective layout, got\"\n            f\" {b_scale.layout}\"\n        )\n    elif m == 128:\n      if b_scale.layout != scales_layout():\n        raise ValueError(\n            f\"Expected B scales to have a M=128 layout, got {b_scale.layout}\"\n        )\n    else:\n      raise AssertionError(\"Should not happen\")\n    if b_scale.shape[0] % 128 or b_scale.shape[0] < n * num_cta:\n      raise ValueError(\n          f\"B scale shape[0] must be a multiple of 128 and >= N={n * num_cta},\"\n          f\" got {b_scale.shape[0]}\"\n      )","sourceCodeStart":454,"sourceCodeEnd":490,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/tcgen05.py#L454-L490","documentation":"The A scale tensor in tcgen05.mma must carry the specific TMEM layout returned by scales_layout(). Any other layout (or None) is rejected because the tensor-core instruction reads scales from fixed TMEM lanes.","triggerScenarios":"Passing a_scale whose .layout attribute is not tcgen05.scales_layout(), e.g. a plain (row-major) layout or a swizzled SMEM layout.","commonSituations":"Building scale memrefs with generic layouts or reusing operand layouts for scales; missing the layout wrapping step in a kernel authoring helper.","solutions":["Construct a_scale with the layout from tcgen05.scales_layout()","Check that a_scale is a TensorMemRefView-like object with .layout set, not a raw memref"],"exampleFix":"# before\na_scale = TensorMemRefView(get_buffer(), (TMEM_ROWS, k_scales), dtype, layout=None)\n# after\na_scale = TensorMemRefView(get_buffer(), (TMEM_ROWS, k_scales), dtype, layout=scales_layout())","handlingStrategy":"validation","validationCode":"assert a_scale.layout == scales_layout(), 'bad A scale layout'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Always construct scale views with scales_layout() from tcgen05","Never reuse operand layouts for scale tensors"],"tags":["gpu","mosaic","tcgen05","layout","tmem"],"backgroundTag":"invalid-memory-layout","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}