fix mx kernel tests (#2614)

vkuzo · web-flow · commit bf5bd5ff31b8 · 2025-07-28T13:16:23.000-04:00
Update

[ghstack-poisoned]
diff --git a/torchao/prototype/mx_formats/kernels.py b/torchao/prototype/mx_formats/kernels.py
@@ -1394,7 +1394,7 @@ def triton_to_mxfp8_dim1_reference(
         scale_e8m0_dim1 = scale_e8m0_dim1.view(torch.float8_e8m0fnu)
         return (
             x_hp_d1_normalized.t(),
-            scale_e8m0_dim1,
+            scale_e8m0_dim1.unsqueeze(-1),
         )
 
     @triton.jit

Original file line number	Diff line number	Diff line change
`@@ -1394,7 +1394,7 @@ def triton_to_mxfp8_dim1_reference(`
`1394`	`1394`	`scale_e8m0_dim1 = scale_e8m0_dim1.view(torch.float8_e8m0fnu)`
`1395`	`1395`	`return (`
`1396`	`1396`	`x_hp_d1_normalized.t(),`
`1397`		`- scale_e8m0_dim1,`
	`1397`	`+ scale_e8m0_dim1.unsqueeze(-1),`
`1398`	`1398`	`)`
`1399`	`1399`
`1400`	`1400`	`@triton.jit`