Porting PyTorch weight to Jax

[ranlucienwang]

Assume we have a PyTorch Model and a Jax model. Is there a framework where you can port PyTorch layer weight to Jax? I might need to implement many models from PyTorch to Jax, and the only way I can think of that can test the correctness of the algorithm is by initializing and then porting the models.

[rwightman]

@ranlucienwang moving to discussions, I fiddled with a jax port of efficientnet a while back (https://github.com/rwightman/efficientnet-jax), it's not too hard to translate weights just need to come up with a scheme that works well with whatever lib you want to use in jax,