autodidax: jit, DeviceArrays, multi-output, pytrees

[mattjj]

The changes here, in no particular order, are roughly:

• Added transpose and broadcast primitives, b/c they're needed by vmap. (We could instead switch to a simplified version of vmap that doesn't need these, but this way it's a bit more faithful to the real jax.vmap.)
• Updated jaxprs, primitives, and transformations to be multi-output, since we need that for jvp (especially jvp-of-jit).
• For the preceding bullet, added pytrees!
• Added jit, with caching and DeviceArrays, as well as jvp and vmap rules for it. That entailed adding an eval_jaxpr, and not much else!
• Fixed a subtle bug where we drop references to Tracers and their memory addresses get reused. It's funny because I had the same bug in omnistaging #3370.
• Added some text about "initial-style" vs "final-style".

$ cloc docs/autodidax.py
       1 text file.
       1 unique file.
       0 files ignored.

github.com/AlDanial/cloc v 1.86  T=0.01 s (71.0 files/s, 117368.4 lines/s)
-------------------------------------------------------------------------------
Language                     files          blank        comment           code
-------------------------------------------------------------------------------
Python                           1            295            460            898
-------------------------------------------------------------------------------

I think we'll be able to get grad in here at under 1000 SLoC, though probably not final-style calls (i.e. custom_jvp) or pmap.

[froystig]

Nice.