[JAX] Remove references to deprecated types DeviceArray and ShardedDe…

…viceArray.

Both are slated for deletion, and neither can be produced at runtime if jax.Array is enabled, which will be the only option as of March 15 (go/jax-array).

PiperOrigin-RevId: 517200975

neural_tangents/_src/batching.py

@@ -49,7 +49,6 @@
 from jax import device_put, devices
 from jax import jit
 from jax import pmap
-from jax.interpreters.pxla import ShardedDeviceArray
 from jax import random
 import jax.numpy as np
 from jax.tree_util import tree_all
@@ -706,9 +705,6 @@ def jit_or_pmap_broadcast(f: Callable, device_count: int = -1) -> Callable:
     def broadcast(arg: np.ndarray) -> np.ndarray:
       if device_count == 0:
         return arg
-      # If the argument has already been sharded, no need to broadcast it.
-      if isinstance(arg, ShardedDeviceArray) and arg.shape[0] == device_count:
-        return arg
       return np.broadcast_to(arg, (device_count,) + arg.shape)
 
     @utils.wraps(f)