Enabling var_mean decomposition #2273
Conversation
github-actions
bot
added
component: api [Python]
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
There was a problem hiding this comment.
Change looks good! Consider adding aten.rsub, aten.rsqrt, aten.sqrt, or any other decompositions which might be preferable/have recent PRs.
For instance, the following can just be replaced with aten.rsub:
TensorRT/py/torch_tensorrt/dynamo/lowering/_decomposition_groups.py
Lines 128 to 129 in 1ad2aeb
| @@ -175,6 +175,7 @@ | |||
| aten.linalg_vector_norm, | |||
| aten.full, | |||
| aten.repeat, | |||
| aten.var_mean, | |||
There was a problem hiding this comment.
@gs-olive you mentioned the rsub. I will add that.
Regarding sqrt and rsqrt, sqrt is not of lowering type so I don't think it should be included here.
For rsqrt, should I include it here? Since rsqrt is already present in py/torch_tensorrt/dynamo/lowering/_decompositions.py and that would take precedence over the enabled ones, right? And we would not need to add it to the disabled ops then.
There was a problem hiding this comment.
That is true, I agree with your points. I think the rsub would be the main one to change; we can leave the others out, and the rsqrt implementation in _decompositions.py will take precedence, yes.
This PR enables the var_mean and rsub decomposition in
torch_enabled_decompositions