PrimTorch: could we get pure core-aten-ops or prims-ops after aot_autograd

The answer to this question is a little complicated because we have not spent enough effort on this part of the UX, but the short answer is that you have the ability to pick which decomps you use, so in principle you can get only aten or prims in your graph by choosing your decomps carefully. How to choose? Well, uh, we should have an api for this but we don’t. We could probably figure this out automatically by analyzing each decomp and seeing if it calls a prim or not. This would be a pretty useful PR if you want to submit it.

2 Likes