I think that Objax is a JAX-based neural network library that is closest to PyTorch. This seems to be a good candidate to have something like PyroModule in NumPyro.
interested in taking this issue :)
Most helpful comment
interested in taking this issue :)