Some recent discussions on “wrapping user-defined JAX functions” made me wonder about using JAX directly as a backend for PyMC, without going through Aesara.
Besides the ability to compile to multiple backends, are there other benefits of having Aesara as a layer between PyMC3 and JAX?
If one is interested in only using JAX as a backend, would there be a benefit in dropping Aesara completely?