Has anyone had a satisfying experience using numpyro with a physical units library?
I see several libraries offering partial / limited support for jax programs, none of which stands out to my eye as likely integrating well with numpyro. I think I’d want units to work their ways through distributions sometimes, in distribution-dependent ways, and I’d maybe even want something to indicate e.g. “this is the log of a value in meters”, which could be exponentiated to get meters again.