Merge pull request jax-ml#24417 from rajasekharporeddy:testbranch1

PiperOrigin-RevId: 688159150

Author: Google-ML-Automation

jax/_src/interpreters/pxla.py

@@ -1868,7 +1868,8 @@ def _raise_warnings_or_errors_for_jit_of_pmap(
          "does not preserve sharded data representations and instead collects "
          "input and output arrays onto a single device. "
          "Consider removing the outer jit unless you know what you're doing. "
-         "See https://github.com/jax-ml/jax/issues/2926.")
+         "See https://github.com/jax-ml/jax/issues/2926. Or "
+         "use jax.experimental.shard_map instead of pmap under jit compilation.")
 
   if nreps > xb.device_count(backend):
     raise ValueError(