Refactoring: rename dtypes.bit_width -> itemsize_bits

justinjfu · Google-ML-Automation · commit 5d06d3a1803d · 2025-11-21T13:35:24.000-08:00
PiperOrigin-RevId: 835341683
diff --git a/tests/mosaic/gpu_test.py b/tests/mosaic/gpu_test.py
@@ -1083,7 +1083,7 @@ def kernel(ctx, lhs, rhs, out, rhs_smem):
           m=m, n=n, dtype=utils.dtype_to_ir_type(out_dtype),
           is_signed=True if dtype == jnp.int8 else None,
       )
-      layout = fa.WGMMA_LAYOUT_8BIT if dtypes.bit_width(dtype) == 8 else fa.WGMMA_LAYOUT
+      layout = fa.WGMMA_LAYOUT_8BIT if dtypes.itemsize_bits(dtype) == 8 else fa.WGMMA_LAYOUT
       lhs_regs = fa.FragmentedArray.load_untiled(
           lhs, layout=layout, optimized=False, is_signed=utils.is_signed(dtype),
       )

Original file line number	Diff line number	Diff line change
`@@ -1083,7 +1083,7 @@ def kernel(ctx, lhs, rhs, out, rhs_smem):`
`1083`	`1083`	`m=m, n=n, dtype=utils.dtype_to_ir_type(out_dtype),`
`1084`	`1084`	`is_signed=True if dtype == jnp.int8 else None,`
`1085`	`1085`	`)`
`1086`		`- layout = fa.WGMMA_LAYOUT_8BIT if dtypes.bit_width(dtype) == 8 else fa.WGMMA_LAYOUT`
	`1086`	`+ layout = fa.WGMMA_LAYOUT_8BIT if dtypes.itemsize_bits(dtype) == 8 else fa.WGMMA_LAYOUT`
`1087`	`1087`	`lhs_regs = fa.FragmentedArray.load_untiled(`
`1088`	`1088`	`lhs, layout=layout, optimized=False, is_signed=utils.is_signed(dtype),`
`1089`	`1089`	`)`