jax

23eaf216 - Make inspect_array_sharding work without mesh context manager too.