jax

8ac19c72 - Fix a faulty soft_pmap rule for axis_index

Fix a faulty soft_pmap rule for axis_index The rule didn't specify the precision for the `np.arange` constant, which caused an accidental dtype promotion in X64 mode. Previously the error has luckicly been hidden behind a coerction that followed `axis_index` in that test, but the new implementation has surfaced it.