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.

jax/lax/lax_parallel.py

@@ -678,7 +678,7 @@ def _axis_index_translation_rule(c, *, axis_name, axis_env, platform):
 def _axis_index_soft_pmap_rule(vals, mapped, chunk_size, *, axis_name):
   assert not vals and not mapped
   idx = axis_index(axis_name)  # type: ignore
-  return idx * chunk_size + np.arange(chunk_size), True
+  return idx * chunk_size + np.arange(chunk_size, dtype=np.int32), True
 
 axis_index_p = core.Primitive('axis_index')
 xla.parallel_translations[axis_index_p] = _axis_index_translation_rule