jax

1e3f7005 - print partial eval stats in jit