[bevenky]

This is also being added to pytorch

https://github.com/pytorch/pytorch/pull/46750

[minimaxir]

I don't think that's the Stage 3 announced in this blog post, but it's def a framework for it.