[TorchAO] A question about the lowering for `to_mx` cast

username: tatwaichong
cooked: <p>I followed the MXTensor prototype (<a href="https://github.com/pytorch/ao/blob/main/torchao/prototype" title="https://github.com/pytorch/ao/blob/main/torchao/prototype" rel="noopener nofollow ugc">https://github.com/pytorch/ao/blob/main/torchao/prototype</a>) to use to_mx() to cast fp32 data to mx (float8_e4m3fn, float8_f8e8m0fnu)</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
import torch.nn as nn
from torchao.prototype.mx_formats.mx_tensor import MXTensor

class example(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, x):
        # high precision to MX, block size defaults to 32
        x_mx = MXTensor.to_mx(x, torch.float8_e4m3fn)
        return x_mx

model = example()

x = torch.randn(32, 32, device='cpu')

output = torch.export.export(model, (x,))
print(output)
</code></pre>
<p>Apparently to_mx() is lowered into a bunch of ATen IRs, as shown</p>
<pre data-code-wrap="python"><code class="lang-python">def forward(self, c_lifted_tensor_0: "f32[32, 32]", x: "f32[32, 32]"):
        reshape: "f32[32, 1, 32]" = torch.ops.aten.reshape.default(x, [32, 1, 32]);  x = None
        abs_1: "f32[32, 1, 32]" = torch.ops.aten.abs.default(reshape)
        amax: "f32[32, 1]" = torch.ops.aten.amax.default(abs_1, [-1]);  abs_1 = None
        unsqueeze: "f32[32, 1, 1]" = torch.ops.aten.unsqueeze.default(amax, -1);  amax = None
        to: "f32[32, 1, 32]" = torch.ops.aten.to.dtype(reshape, torch.float32);  reshape = None
        to_1: "f32[32, 1, 1]" = torch.ops.aten.to.dtype(unsqueeze, torch.float32);  unsqueeze = None
        view: "i32[32, 1, 1]" = torch.ops.aten.view.dtype(to_1, torch.int32)
        bitwise_right_shift: "i32[32, 1, 1]" = torch.ops.aten.bitwise_right_shift.Tensor_Scalar(view, 23);  view = None
        and_1: "i32[32, 1, 1]" = torch.ops.aten.__and__.Scalar(bitwise_right_shift, 255);  bitwise_right_shift = None
        sub: "i32[32, 1, 1]" = torch.ops.aten.sub.Tensor(and_1, 127);  and_1 = None
        sub_1: "i32[32, 1, 1]" = torch.ops.aten.sub.Tensor(sub, 8);  sub = None
        clamp: "i32[32, 1, 1]" = torch.ops.aten.clamp.default(sub_1, -127, 128);  sub_1 = None
        add: "i32[32, 1, 1]" = torch.ops.aten.add.Tensor(clamp, 127);  clamp = None
        to_2: "u8[32, 1, 1]" = torch.ops.aten.to.dtype(add, torch.uint8);  add = None
        isnan: "b8[32, 1, 1]" = torch.ops.aten.isnan.default(to_1);  to_1 = None
        where: "u8[32, 1, 1]" = torch.ops.aten.where.ScalarSelf(isnan, 255, to_2);  isnan = to_2 = None
        to_3: "i32[32, 1, 1]" = torch.ops.aten.to.dtype(where, torch.int32)
        bitwise_left_shift: "i32[32, 1, 1]" = torch.ops.aten.bitwise_left_shift.Tensor_Scalar(to_3, 23);  to_3 = None
        view_1: "f32[32, 1, 1]" = torch.ops.aten.view.dtype(bitwise_left_shift, torch.float32);  bitwise_left_shift = None
        clamp_1: "f32[32, 1, 1]" = torch.ops.aten.clamp.default(view_1, 1.1754943508222875e-38);  view_1 = None
        div: "f32[32, 1, 32]" = torch.ops.aten.div.Tensor(to, clamp_1);  to = clamp_1 = None
        to_4: "f8e4m3fn[32, 1, 32]" = torch.ops.aten.to.dtype(div, torch.float8_e4m3fn);  div = None
        reshape_1: "f8e4m3fn[32, 32]" = torch.ops.aten.reshape.default(to_4, [32, 32]);  to_4 = reshape_1 = None
        view_2: "f8e8m0fnu[32, 1, 1]" = torch.ops.aten.view.dtype(where, torch.float8_e8m0fnu);  where = None
        squeeze: "f8e8m0fnu[32, 1]" = torch.ops.aten.squeeze.dim(view_2, -1);  view_2 = squeeze = None
        return (c_lifted_tensor_0,)
</code></pre>
<p>I wonder is this fixed op decomposition final version? I expect different hardware backends might do casting differently — for example different rounding or quantization approach. Do you think to_mx() will eventually get its own ATen op? Or will there be a way for backends to control how to_mx() lowering works?</p>