jax

a2ac234f - Streamline some core.concrete_aval compute paths