num replicas not correctly detected with `lax.switch` / `lax.cond`
Thanks to @LenaMartens for the repro!

```python
import os
os.environ['XLA_FLAGS'] = '--xla_force_host_platform_device_count=32'

import jax

def identity(x):
  return x

def cond_of_pmap(x):
  y = jax.lax.cond(True, lambda x: jax.pmap(identity)(x), lambda x: jax.pmap(identity)(x), x)
  return y

cond_of_pmap(jax.numpy.zeros((8, 2)))
```

```
AssertionError                            Traceback (most recent call last)
google3/third_party/py/jax/interpreters/xla.py in axis_groups(axis_env, name)
    533   mesh_axes = tuple(unsafe_map(partial(axis_read, axis_env), name))
    534   trailing_size, ragged = divmod(axis_env.nreps, prod(axis_env.sizes))
--> 535   assert not ragged
    536   mesh_spec = axis_env.sizes + (trailing_size,)
    537   return _axis_groups(mesh_spec, mesh_axes)

AssertionError: 
```

The issue was in `core.traverse_jaxpr_params`, which assumed that all subjaxprs appear as values in the params dict, while `lax.switch` (aka `lax.cond`) actually has a param value which is a _tuple_ of jaxprs. Those tuple elements were being ignored!

I'm going to make a quick fix which just flattens params to one level, but I think a better approach here would be to have a primitive-keyed table of rules for traversing subjaxprs.
diff --git a/jax/core.py b/jax/core.py
index 115b0e691..5c25057e4 100644
--- a/jax/core.py
+++ b/jax/core.py
@@ -307,11 +307,13 @@ def extract_call_jaxpr(
     return (params["call_jaxpr"], new_params)
 
 
+# TODO(mattjj): replace this approach with a primitive-keyed table of rules
 def traverse_jaxpr_params(f, params):
   """Applies f to each jaxpr parameter and returns a tuple of returned values."""
-  return {name: f(param)
+  return {name: f(p)
           for name, param in params.items()
-          if type(param) in (Jaxpr, ClosedJaxpr)}
+          for p in (param if isinstance(param, (tuple, list)) else [param])
+          if type(p) in (Jaxpr, ClosedJaxpr)}
 
 
 def eval_jaxpr(jaxpr: Jaxpr, consts, *args):
diff --git a/tests/pmap_test.py b/tests/pmap_test.py
index 385b539c2..0bc6827ea 100644
--- a/tests/pmap_test.py
+++ b/tests/pmap_test.py
@@ -1701,6 +1701,17 @@ class PmapTest(jtu.JaxTestCase):
     out_dtype = func(unused_arg).dtype
     self.assertEqual(out_dtype, dtype)
 
+  def test_num_replicas_with_switch(self):
+    # https://github.com/google/jax/issues/7411
+    def identity(x):
+      return x
+
+    def cond_of_pmap(x):
+      y = lax.cond(True, jax.pmap(identity), jax.pmap(identity), x)
+      return y
+
+    cond_of_pmap(jnp.zeros((xla_bridge.device_count(), 2)))
+
 
 class VmapOfPmapTest(jtu.JaxTestCase):
 
['tests/pmap_test.py', 'jax/core.py']