Repository: google/jax
Issue: jnp.arange is broken inside shmap unless using jit

### Description

This may be expected, but I found the error message quite confusing:
```python
from functools import partial
import numpy as np
import chex
import jax
import jax.numpy as jnp
from jax.sharding import Mesh, PartitionSpec as P
from jax.experimental import mesh_utils
from jax.experimental.shard_map import shard_map

chex.set_n_cpu_devices(8)
devices = mesh_utils.create_device_mesh((8,))
mesh = Mesh(devices, axis_names=('i',))

# uncomment this to make it work!
# @jax.jit
@partial(shard_map, mesh=mesh, in_specs=(P('i'),), out_specs=P('i'))
def f(a_block):
  i = jnp.arange(a_block.shape[0])
  return i + a_block

f(np.zeros(32))
```
Raises:
```
XlaRuntimeError                           Traceback (most recent call last)
[<ipython-input-2-9c0c4e291861>](https://localhost:8080/#) in <cell line: 20>()
     18   return i + a_block
     19 
---> 20 f(np.zeros(32))

    [... skipping hidden 5 frame]

2 frames
[<ipython-input-2-9c0c4e291861>](https://localhost:8080/#) in f(a_block)
     15 @partial(shard_map, mesh=mesh, in_specs=(P('i'),), out_specs=P('i'))
     16 def f(a_block):
---> 17   i = jnp.arange(a_block.shape[0])
     18   return i + a_block
     19 

[/usr/local/lib/python3.9/dist-packages/jax/_src/numpy/lax_numpy.py](https://localhost:8080/#) in arange(start, stop, step, dtype)
   2243       ceil_ = ufuncs.ceil if isinstance(start, core.Tracer) else np.ceil
   2244       start = ceil_(start).astype(int)  # type: ignore
-> 2245     return lax.iota(dtype, start)
   2246   else:
   2247     if step is None and start == 0 and stop is not None:

    [... skipping hidden 16 frame]

[/usr/local/lib/python3.9/dist-packages/jax/_src/dispatch.py](https://localhost:8080/#) in backend_compile(backend, built_c, options, host_callbacks)
    469   # TODO(sharadmv): remove this fallback when all backends allow `compile`
    470   # to take in `host_callbacks`
--> 471   return backend.compile(built_c, compile_options=options)
    472 
    473 _ir_dump_counter = itertools.count()

XlaRuntimeError: INTERNAL: RET_CHECK failure (external/xla/xla/hlo/ir/hlo_instruction.cc:2828) ShapeUtil::CompatibleIgnoringFpPrecision(shape(), new_producer->shape()) s32[8,4]{1,0} is not compatible with s32[1,4]{1,0}
```

### What jax/jaxlib version are you using?

_No response_

### Which accelerator(s) are you using?

_No response_

### Additional system info

_No response_

### NVIDIA GPU info

_No response_

Changed files:
jax/_src/dispatch.py, tests/shard_map_test.py

diff --git a/jax/_src/dispatch.py b/jax/_src/dispatch.py
index bdd2a7325..cd6c36343 100644
--- a/jax/_src/dispatch.py
+++ b/jax/_src/dispatch.py
@@ -320,7 +320,7 @@ def jaxpr_shardings(
         ndmin = max(names) + 1 if names else 0
         return PartitionSpec(*(names.get(i) for i in range(ndmin)))
       yield from ((NamedSharding(eqn.params['mesh'], _names_to_pspec(names)), source_info)
-                  for names in eqn.params['in_names'])
+                  for names in [*eqn.params['in_names'], *eqn.params['out_names']])
   for subjaxpr in core.subjaxprs(jaxpr):
     yield from jaxpr_shardings(subjaxpr)
 
diff --git a/tests/shard_map_test.py b/tests/shard_map_test.py
index ccdec63b8..7b9a4118e 100644
--- a/tests/shard_map_test.py
+++ b/tests/shard_map_test.py
@@ -684,6 +684,23 @@ class ShardMapTest(jtu.JaxTestCase):
     with self.assertRaisesRegex(NotImplementedError, 'axis_index'):
       g(x)
 
+  def test_jaxpr_shardings_with_no_outputs(self):
+    # https://github.com/google/jax/issues/15385
+    mesh = jtu.create_global_mesh((4,), ('i',))
+
+    @jax.jit
+    @partial(shard_map, mesh=mesh, in_specs=(), out_specs=P('i'))
+    def f():
+      return jax.lax.iota(jnp.dtype('int32'), 4)
+    f()  # don't crash
+
+    @partial(shard_map, mesh=mesh, in_specs=(P('i'),), out_specs=P('i'))
+    def g(a_block):
+      i = jnp.arange(a_block.shape[0])
+      return i + a_block
+
+    g(np.arange(32))  # don't crash
+
 
 class FunSpec(NamedTuple):
   name: str