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_


Discussion:
Also fails if the body
```python
def f(a_block):
  i = jnp.arange(a_block.shape[0])
  return i + a_block
```

is replaced with
```python
def f(a_block):
  jnp.arange(4)  # or even jax.lax.iota(jnp.dtype('int32'), 4)
```
Yet replacing `jnp.arange(4)` with `jnp.zeros(4)` or similar works fine... I wonder if this is an XLA bug.

Does it look like there's anything wrong with this HLO?

```
HloModule jit__unnamed_wrapped_function_, entry_computation_layout={()->f32[8,4]{1,0}}

ENTRY main.5 {
  iota.1 = f32[4]{0} iota(), iota_dimension=0
  reshape.2 = f32[1,4]{1,0} reshape(iota.1)
  custom-call.3 = f32[1,4]{1,0} custom-call(reshape.2), custom_call_target="Sharding", sharding={manual}
  ROOT custom-call.4 = f32[8,4]{1,0} custom-call(custom-call.3), custom_call_target="SPMDShardToFullShape", sharding={
devices=[8,1]0,1,2,3,4,5,6,7}
}
```
Interestingly, `jax.jit` doesn't always fix it: it's only a `jax.jit` _and at least one argument passed into the shmap_ which fixes it, e.g. this fails:

```python
@jax.jit
@partial(shard_map, mesh=mesh, in_specs=(), out_specs=P('i'))
def f():
  return jax.lax.iota(jnp.dtype('int32'), 4)
f()
```

I think there's a bug in XLA "manual spmd mode" with zero-input functions. (The code in the OP hits this case because eager shmap basically means op-by-op shmap, so we hit the iota/arange by itself and hence build a computation with no inputs, whereas adding `jax.jit` fixes the code in the OP because it happens to have an input argument.)
I filed internal bug b/277140784
Oops, this is a JAX bug after all; somehow in the problematic code we're telling XLA to compile with `num_partitions=1`. (Thanks @tongfei-guo for patiently correcting me!)

Resolving patch:
diff --git a/jax/_src/dispatch.py b/jax/_src/dispatch.py
--- a/jax/_src/dispatch.py
+++ b/jax/_src/dispatch.py
@@ -320,7 +320,7 @@ def _names_to_pspec(names):
         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)
 


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
@@ -684,6 +684,23 @@ def foo(x):
     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