I would start here: https://jax.readthedocs.io/en/latest/notebooks/Writing_custom_interpreters_in_Jax.html https://jax.readthedocs.io/en/latest/notebooks/Writin

I would start here: https://jax.readthedocs.io/en/latest/notebooks/Writing_custom_interpreters_in_Jax.html https://jax.readthedocs.io/en/latest/notebooks/Writing_custo... Note that JAX enforces certain limitations, so you should be careful when considering JAX to be "the perfect language" - in general I don't think this is true. It's quite good at what it's designed for.