[btwillard]

I'm the author of the JAX backend in Theano, and, no, this didn't have anything to do with NumPyro--especially since I neither use nor know very much about NumPyro.

It was just a quick demonstration of how easily one can use Theano as a generalized graph "front-end", while also preserving its more unique and programmable symbolic optimization capabilities. JAX was one of a few "backends" I considered, and, due to the JAX Python library, it also looked like the most straightforward one to implement first.