bug: jax v0.4.36

[nstarman]

quax is calling jax.core.new_main, which no longer exists (or has been moved maybe?) in jax v0.4.36
EDIT:
I think this is related to the changes made to jax's tracers. I can still find a new_main in jax's autodidax.py example, but nowhere else.

[nstarman]

I traced it to jax-ml/jax@c36e1f7#diff-713d667c14b1dd075ec025dbd1cd8ea222e42c8b4c094ae05373d4d05b276dc5L1223, which removes new_main.

[patrick-kidger]

Yup, JAX have changed how tracing works. See also jax-ml/jax#25289

If you have the time and interest I'd be very happy to take a PR fixing things up here!

[nstarman]

I was poking around, but I really don't know enough about the internals of Jax's tracing machinery.
They don't obviously document a replacement function, nor alternative. 😓
I'll open an Issue asking.

[nstarman]

I opened jax-ml/jax#25372 (comment) and was told that there will be some examples developed for autodidax, and that the ad system had to make changes that are likely similar to those needed in quax — jax-ml/jax#25372 (comment)

[nstarman]

I think #37 will close this!