fix: current failing jax tests for jax>=0.4.36

[pfackeldey]

This PR fixes the current failing JAX tests since jax version 0.4.36. It adds a flag to restore old (jax<0.4.36) tracing behavior, see more at: https://jax.readthedocs.io/en/latest/changelog.html#jax-0-4-36-dec-5-2024

At some point we likely want to address this properly, but for now this should make our jax tests pass again.

[ianna]

@pfackeldey - Thanks for prompt fix! Looks good to me.

[jpivarski]

Great, this fixes the tests! But it's a good example of how hard it is to work with JAX: jax.config.update is not in our control. It only takes effect if it is the first thing that's called after importing JAX. After that time, it's silently ignored!

[jpivarski]

About the WASM failure: that one would be fixed by #3333. In fact, we have a cross-dependence and can't merge either one first without errors. I'll force-merge #3333 and then you can pick that up here, then we'll merge this one and the other PRs can pick up on that.

[martindurant]

(#3333 is now merged, if you were waiting)

[henryiii]

The jobs in #3333 are still required, so non-admins can no longer merge. Why was that job removed?

[jpivarski]

The WASM-building job in docs.yml was removed because it was building a WASM version of awkward-cpp for the documentation, but the documentation no longer includes an embedded JupyterLite iframe. (Instead, the "try it" button points to a raw-Pyodide page that gets awkward-cpp from the web.)

I don't know how #3333 passed if that's still required. I don't remember it waiting for that reason. But I'll remove the requirement now.