[shard-map] Check failed using unsafe_rbg_key
Based on code from @jheek :

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

from functools import partial

import numpy as np

import jax
import jax.numpy as jnp
from jax import random
from jax.experimental.shard_map import shard_map, PartitionSpec as P, Mesh

from flax import linen

jax.config.update('jax_enable_custom_prng', True)  # doesn't crash if we exclude this line!

dense = linen.Dense(128)

mesh = Mesh(np.array(jax.devices()).reshape(2, 4), ('model', 'data'))

@jax.jit
def f(rng):
  @partial(shard_map, mesh=mesh, in_specs=P("model"), out_specs=P("model"), check_rep=False)
  def g(rng):
    rng = rng[0]
    return dense.init(rng, jnp.zeros((8, 128)))

  return g(random.split(rng, 4))

f(random.PRNGKey(0)) # ==> crashes VM on ASSERT(IsManual(sharding))
```
diff --git a/jax/_src/interpreters/mlir.py b/jax/_src/interpreters/mlir.py
index 8dde30217..d6d0196ef 100644
--- a/jax/_src/interpreters/mlir.py
+++ b/jax/_src/interpreters/mlir.py
@@ -1500,7 +1500,7 @@ def _wrap_with_spmd_op(name: str,
   else:
     backend_config = ""
   result_type = aval_to_ir_type(aval_out)
-  out_shape = aval_out.shape  # type: ignore
+  out_shape = core.physical_aval(aval_out).shape  # type: ignore
   if core.is_constant_shape(out_shape):
     result_shapes = None
   else:
diff --git a/jax/_src/prng.py b/jax/_src/prng.py
index d23283eec..a62307c80 100644
--- a/jax/_src/prng.py
+++ b/jax/_src/prng.py
@@ -470,7 +470,8 @@ class KeyTyRules:
       op_sharding_proto = op_sharding_proto.to_proto()
     new_op_sharding = op_sharding_proto.clone()
     tad = list(new_op_sharding.tile_assignment_dimensions)
-    tad.extend([1] * len(key_shape))
+    suffix = [tad.pop()] if op_sharding_proto.replicate_on_last_tile_dim else []
+    tad.extend([1] * len(key_shape) + suffix)
     new_op_sharding.tile_assignment_dimensions = tad
     return xc.HloSharding.from_proto(new_op_sharding)
 
diff --git a/tests/shard_map_test.py b/tests/shard_map_test.py
index 6b8aaf1e3..5f28ce454 100644
--- a/tests/shard_map_test.py
+++ b/tests/shard_map_test.py
@@ -740,6 +740,20 @@ class ShardMapTest(jtu.JaxTestCase):
     # error!
     jax.jit(g)(x)  # doesn't crash
 
+  def test_key_array_with_replicated_last_tile_dim(self):
+    # See https://github.com/google/jax/issues/16137
+
+    mesh = jtu.create_global_mesh((2, 4), ('i', 'j'))
+
+    def f(rng):
+      @partial(shard_map, mesh=mesh, in_specs=P('i'), out_specs=P('i'),
+               check_rep=False)
+      def g(rng):
+        return jnp.array([jax.random.normal(rng[0])])
+      return g(jax.random.split(rng, 4))
+
+    jax.jit(f)(jax.random.key(0))  # doesn't crash
+
   # same method appears in api_test.py:DCETest
   # TODO(mattjj): consider moving this method to be a helper in jtu
   def assert_dce_result(self, jaxpr: core.Jaxpr, used_outputs: List[bool],
['jax/_src/interpreters/mlir.py', 'jax/_src/prng.py', 'tests/shard_map_test.py']