jax

e0d1b375 - Delete dead axis_index code

Delete dead axis_index code The primitive was moved to `lax_parallel.py` some time ago, so the one in `core` should no longer be used. This is probably a result of a botched rebase.