Repository: google/jax

Issue:
`shard_map` throws a `NotImplementedError: No replication rule for erf_inv`
### Description

`shard_map` throws a `NotImplementedError: No replication rule for erf_inv`

Repro:
```
# Create a mesh and annotate the axis with a name.
device_mesh = mesh_utils.create_device_mesh((8,))
mesh = Mesh(devices=device_mesh, axis_names=('data',))
data_sharding = NamedSharding(mesh, PartitionSpec('data',))

class Model(nn.Module):
  @nn.compact
  def __call__(self, x, add_noise):
    x = nn.Dense(1)(x)
    # use jnp.where for control flow; for more details see: https://jax.readthedocs.io/en/latest/errors.html#jax.errors.TracerBoolConversionError
    return jnp.where(
      add_noise, x + jax.random.normal(self.make_rng('params'), x.shape), x
    )

module = Model()
init_rng, apply_rng = jax.random.split(jax.random.key(0))
# get 8 different rng's that will be used by the 8 devices when doing forward inference
apply_rng = jax.random.split(apply_rng, 8)
x = jnp.ones((8, 1))
variables = module.init(init_rng, x, False)

def forward(variables, x, add_noise, rng_key_batch):
  # rng_key_batch is a batch of size 1 containing 1 PRNG key
  # index slice into the rng_key_batch to access the PRNG key
  return module.apply(
    variables, x, add_noise, rngs={'params': rng_key_batch[0]}
  )

# define partition specifications
data_pspec = PartitionSpec('data')
no_pspec = PartitionSpec()

# shard the inputs x and rng keys across devices
# replicate the variables and add_noise boolean across devices
# shard the output across devices
# set `check_rep=False`, otherwise we get a Jax error
shmap_forward = shard_map(
  forward,
  mesh=mesh,
  in_specs=(no_pspec, data_pspec, no_pspec, data_pspec),
  out_specs=data_pspec,
)
out = shmap_forward(variables, x, True, apply_rng)
```
If we add `check_rep=False`, we don't get an error:
```
shmap_forward = shard_map(
  forward,
  mesh=mesh,
  in_specs=(no_pspec, data_pspec, no_pspec, data_pspec),
  out_specs=data_pspec,
  check_rep=False,
)
```

### What jax/jaxlib version are you using?

0.4.21

### Which accelerator(s) are you using?

CPU

### Additional system info?

_No response_

### NVIDIA GPU info

_No response_


Discussion:

Resolving patch:
diff --git a/jax/experimental/shard_map.py b/jax/experimental/shard_map.py
--- a/jax/experimental/shard_map.py
+++ b/jax/experimental/shard_map.py
@@ -48,7 +48,7 @@
 from jax._src.api import _shared_code_pmap, _prepare_pmap
 from jax._src.lax import (lax, parallel as lax_parallel, slicing,
                           windowed_reductions, convolution, fft, linalg,
-                          control_flow)
+                          special, control_flow)
 from jax._src.util import (HashableFunction, HashablePartial, unzip2, unzip3,
                            as_hashable_function, memoize, partition_list,
                            merge_lists, split_list, subs_list2,
@@ -909,9 +909,10 @@ def _standard_collective_rewrite(prim, mesh, in_rep, x, axis_name, **params):
 
 for o in it.chain(lax.__dict__.values(), slicing.__dict__.values(),
                   windowed_reductions.__dict__.values(),
-                  convolution.__dict__.values(), fft.__dict__.values(),
-                  linalg.__dict__.values(), ops.__dict__.values(),
-                  ad_util.__dict__.values(), prng.__dict__.values()):
+                  special.__dict__.values(), convolution.__dict__.values(),
+                  fft.__dict__.values(), linalg.__dict__.values(),
+                  ops.__dict__.values(), ad_util.__dict__.values(),
+                  prng.__dict__.values()):
   if isinstance(o, core.Primitive):
     register_standard_check(o)
     register_standard_rewrite(o)


Test patch:
diff --git a/tests/shard_map_test.py b/tests/shard_map_test.py
--- a/tests/shard_map_test.py
+++ b/tests/shard_map_test.py
@@ -1298,6 +1298,12 @@ def test_custom_jvp_inside_jit(self):
                   mesh=mesh, in_specs=P('batch'),
                   out_specs=P('batch'))(jnp.arange(16.))  # don't crash
 
+  def test_random_normal_rules(self):
+    mesh = jtu.create_global_mesh((4,), ('i',))
+    keys = jax.random.split(jax.random.key(0), 4)
+    shard_map(lambda k: jax.random.normal(k[0], (1,)),
+              mesh=mesh, in_specs=P('i'), out_specs=P('i'))(keys)  # don't crash
+
 
 class FunSpec(NamedTuple):
   name: str