Check for jax.Sharding's number of devices instead of `py_array.num…

…_shards` which looks at IFRT sharding's num_devices to check against `global_devices` and deciding whether to fall back to python shard_arg.

This is because IFRT sharding's `num_shards` method is busted. It doesn't return the global shards (in some cases) which leads to JAX program unnecessarily falling back to python.

PiperOrigin-RevId: 673067095

jax/_src/test_util.py

@@ -39,6 +39,7 @@
 import jax
 from jax import lax
 from jax._src import api
+from jax._src import array
 from jax._src import config
 from jax._src import core
 from jax._src import dispatch
@@ -383,6 +384,25 @@ def mlir_lower_and_count(*args, **kwargs):
     mlir.lower_jaxpr_to_module = mlir_lower
 
 
+@contextmanager
+def count_jax_array_shard_arg_calls():
+  # No need to clear any caches since we generally jit and pmap fresh callables
+  # in tests.
+
+  array_shard_arg = array._array_shard_arg
+  count = [0]
+
+  def array_shard_arg_and_count(*args, **kwargs):
+    count[0] += 1
+    return array_shard_arg(*args, **kwargs)
+
+  pxla.shard_arg_handlers[array.ArrayImpl] = array_shard_arg_and_count
+  try:
+    yield count
+  finally:
+    pxla.shard_arg_handlers[array.ArrayImpl] = array_shard_arg
+
+
 @contextmanager
 def count_jit_compilation_cache_miss():
   # No need to clear any caches since we generally jit and pmap fresh callables
@@ -1965,7 +1985,7 @@ def arcsin(self, x):
       # On branch cut, mpmath.mp.asin returns different value compared
       # to mpmath.fp.asin and numpy.arcsin (see
       # mpmath/mpmath#786). The following if-block ensures
-      # compatibiliy with numpy.arcsin.
+      # compatibility with numpy.arcsin.
       if x.real > 1 and x.imag == 0:
         return ctx.asin(x).conjugate()
 
@@ -1997,7 +2017,7 @@ def arccos(self, x):
         return ctx.make_mpc((real._mpf_, (-sign_imag * inf)._mpf_))
       # On branch cut, mpmath.mp.acos returns different value
       # compared to mpmath.fp.acos and numpy.arccos. The
-      # following if-block ensures compatibiliy with
+      # following if-block ensures compatibility with
       # numpy.arccos.
       if x.imag == 0 and x.real > 1:
         return -ctx.acos(x)
@@ -2026,7 +2046,7 @@ def arcsinh(self, x):
       # On branch cut, mpmath.mp.asinh returns different value
       # compared to mpmath.fp.asinh and numpy.arcsinh (see
       # mpmath/mpmath#786).  The following if-block ensures
-      # compatibiliy with numpy.arcsinh.
+      # compatibility with numpy.arcsinh.
       if x.real == 0 and x.imag < -1:
         return (-ctx.asinh(x)).conjugate()
     return ctx.asinh(x)