Any chance to preserve some ops while decomposing PT2E model?

torch.compile FX (Functional Transformations)
1

Hi team,

I wonder is there any way to preserve some type of ops while doing decomposition ep.run_decompositions()?

For example,

aten.upsample_nearest2d.vec   (div_5, None, [2.0, 2.0])

will convert to:

opcode         name          target                   args                                                   
-------------  ------------  -----------------------  -------------------------------------------------------
call_function  arange_0      aten.arange.start_step   (0, 64)                                                
call_function  add_9         aten.add.Tensor          (arange_0, 0.0)                                        
call_function  mul_12        aten.mul.Tensor          (add_9, 0.5)                                           
call_function  _to_copy_23   aten._to_copy.default    (mul_12,)                                              
call_function  unsqueeze_0   aten.unsqueeze.default   (_to_copy_23, -1)                                      
call_function  arange_1      aten.arange.start_step   (0, 64)                                                
call_function  add_10        aten.add.Tensor          (arange_1, 0.0)                                        
call_function  mul_13        aten.mul.Tensor          (add_10, 0.5)                                          
call_function  _to_copy_24   aten._to_copy.default    (mul_13,)                                              
call_function  index_0       aten.index.Tensor        (_to_copy_22, [None, None, unsqueeze_0, _to_copy_24])

But I wish to keep aten.upsample_nearest2d.vec. Thank you.

Code snippet:

model = ...
sample_inputs = ...
ep = torch.export.export(model, sample_inputs)
ep = ep.run_decompositions()
2

You should be able to preserve upsample_nearest2d.vec now, as of a few weeks ago (run_decompositions() shouldn’t secretly decompose a few ops anymore). Can you try again on a nightly?

1 Like