How to understand the architecture design of Jax？

[south-ocean]

There is relatively little information about Jax at the moment, only the source code and Tutarial.
I want to understand the execution logic of jax and try to contribute code.
As with all AI frameworks, jax has both static and dynamic graph modes, so are the bind rules the same for static and dynamic graph modes? Is the dynamic graph directly through pritimive binding to the XLA hlo, and the static graph through primitive being converted to JAXPR and then binding to the XLA hlo? If so, i see some code about impl_rules[cos_p] = lambda x: [np.cos(x)], so th isn‘t the executation rule about cos_p to numpy and so, so can anyone help me explain it, and does have some resource can let me read the jax source code more clearly, thanks.

[jakevdp]

Hi - thanks for the question! The best resource we have for that currently is probably autodidax, which is a simplified re-implementation of the core features of JAX built from the ground up. It's not exactly like the real JAX code (as you noted, in earlier sections it lowers cos and other operations to numpy rather than XLA for simplicity) but it covers many of the core concepts behind JAX's computational model.

[south-ocean]

Thanks for the answer, as you said above, it confuses me a bit, as the jax docs say, jax rewrites the gpu version of numpy and keeps the api the same, if you lower directly to the official numpy (cpu version), I think this conflicts with the above mentioned and can't get the gpu benefits. So isn't the numpy version of jax implemented based on the backend of xla? Or is the official numpy version as an implementation of the cpu version?

[jakevdp]

Jax doesn’t lower to numpy operations. The code you quoted is from autodidax, which is a tutorial meant to be similar to JAX itself, but with simplifications to make the overall structure easier to follow.

[south-ocean]

I see. Thanks.