PyTorch port?

[derpyplops]

Any plans to do this? I might be interested to try working on this if not

[modricwang]

Just do it!

[Sohl-Dickstein]

We've gotten a lot of messages asking for a PyTorch port. We have no current plans to make one ourselves. We would love + be grateful for someone to take it on!!

There was one other message from someone saying they might make a port (as opposed to just wishing there was one): https://mastodon.social/@Sdatkinson/109378948264886553

One quick high level pieces of advice:

• I would choose a simple, multi-tensor, optimization problem, and make sure the JAX and PyTorch versions generate the same optimization trajectory on both. When porting a NN there are a lot of ways to implement a subtly different function that still somewhat works.

Thank you! And let us know how it goes!!