dequantize_per_channel triton graph

gistfile1.txt

 ===== __compiled_fn_12 =====
 <eval_with_key>.135 class GraphModule(torch.nn.Module):
    def forward(self, L_int_repr_ : torch.Tensor, L_scales_ : torch.Tensor, L_zero_points_ : torch.Tensor):
        l_int_repr_ = L_int_repr_
        l_scales_ = L_scales_
        l_zero_points_ = L_zero_points_
        
        # File: /fsx/users/cdhernandez/protoquant/ao_experimental/quant_primitives.py:134, code: y = int_repr.transpose(-1, axis)
        transpose = l_int_repr_.transpose(-1, 2);  l_int_repr_ = None
        
        # File: /fsx/users/cdhernandez/protoquant/ao_experimental/quant_primitives.py:135, code: y = y.to(out_dtype)
        to = transpose.to(torch.float16);  transpose = None
        
        # File: /fsx/users/cdhernandez/protoquant/ao_experimental/quant_primitives.py:136, code: y = y - zero_points
        sub = to - l_zero_points_;  to = l_zero_points_ = None
        
        # File: /fsx/users/cdhernandez/protoquant/ao_experimental/quant_primitives.py:137, code: y = y * scales.to(out_dtype)
        to_1 = l_scales_.to(torch.float16);  l_scales_ = None
        mul = sub * to_1;  sub = to_1 = None
        
        # File: /fsx/users/cdhernandez/protoquant/ao_experimental/quant_primitives.py:138, code: y = y.transpose(-1, axis)
        transpose_1 = mul.transpose(-1, 2);  mul = None
        return (transpose_1,)
        

[2023-05-10 01:03:27,673] torch._dynamo.output_graph.__graph: [DEBUG] TRACED GRAPH
 __compiled_fn_12 <eval_with_key>.135 opcode         name            target                   args                        kwargs
-------------  --------------  -----------------------  --------------------------  --------
placeholder    l_int_repr_     L_int_repr_              ()                          {}
placeholder    l_scales_       L_scales_                ()                          {}
placeholder    l_zero_points_  L_zero_points_           ()                          {}
call_method    transpose       transpose                (l_int_repr_, -1, 2)        {}
call_method    to              to                       (transpose, torch.float16)  {}
call_function  sub             <built-in function sub>  (to, l_zero_points_)        {}
call_method    to_1            to                       (l_scales_, torch.float16)  {}
call_function  mul             <built-in function mul>  (sub, to_1)                 {}
call_method    transpose_1     transpose                (mul, -1, 2)                {}
output         output          output                   ((transpose_1,),)           {}