Repository: google/jax

Issue:
jnp.roll should use simpler approach for static shifts
Current implementation:
```python
>>> jax.make_jaxpr(lambda x: jnp.roll(x, 2))(x)
{ lambda ; a:i32[10]. let
    b:i32[10] = pjit[
      jaxpr={ lambda ; c:i32[10] d:i32[]. let
          e:i32[10] = pjit[
            jaxpr={ lambda ; f:i32[10] g:i32[]. let
                h:i32[1] = broadcast_in_dim[broadcast_dimensions=() shape=(1,)] g
                i:i32[1] = slice[
                  limit_indices=(1,)
                  start_indices=(0,)
                  strides=(1,)
                ] h
                j:i32[] = squeeze[dimensions=(0,)] i
                k:i32[] = convert_element_type[new_dtype=int32 weak_type=False] j
                l:i32[] = max 10 1
                m:i32[] = pjit[
                  jaxpr={ lambda ; n:i32[] o:i32[]. let
                      p:bool[] = eq o 0
                      q:i32[] = pjit[
                        jaxpr={ lambda ; r:bool[] s:i32[] t:i32[]. let
                            u:i32[] = select_n r t s
                          in (u,) }
                        name=_where
                      ] p 1 o
                      v:i32[] = rem n q
                      w:bool[] = ne v 0
                      x:bool[] = lt v 0
                      y:bool[] = lt q 0
                      z:bool[] = ne x y
                      ba:bool[] = and z w
                      bb:i32[] = add v q
                      bc:i32[] = select_n ba v bb
                    in (bc,) }
                  name=remainder
                ] k l
                bd:i32[20] = concatenate[dimension=0] f f
                be:i32[] = sub 10 m
                bf:bool[] = lt be 0
                bg:i32[] = add be 20
                bh:i32[] = select_n bf be bg
                bi:i32[10] = dynamic_slice[slice_sizes=(10,)] bd bh
              in (bi,) }
            name=_roll
          ] c d
        in (e,) }
      name=_roll
    ] a 2
  in (b,) }
```
A more streamlined implementation might look something like this:
```python
>>> jax.make_jaxpr(lambda x: jnp.concatenate([x[-2:], x[:-2]]))(x)
{ lambda ; a:i32[10]. let
    b:i32[2] = dynamic_slice[slice_sizes=(2,)] a 8
    c:i32[8] = dynamic_slice[slice_sizes=(8,)] a 0
    d:i32[10] = concatenate[dimension=0] b c
  in (d,) }
```


Discussion:
Thinking about it, maybe we should just compute the `roll` via a single `gather` over appropriate indices. I think that would be a better implementation in all cases
An important consideration here is what can be lowered into efficient distributed computation by XLA SPMD.

If you look at the SPMD code generated for the existing (dynamic slice) or gather implementations of roll, they are implemented using all-to-all communication (all reduce) rather than nearest neighbor communication (collective permute):
```python
import jax
import jax.numpy as jnp
import chex
import numpy as np

chex.set_n_cpu_devices(4)

def roll_existing(x):
  return jnp.roll(x, +1)

def roll_concat(x):
  return jnp.concatenate([x[-1:], x[:-1]])

def roll_gather(x):
  indices = np.concatenate([[x.size - 1], np.arange(x.size - 1)])
  return x.at[indices].get(unique_indices=True)

np.testing.assert_array_equal(roll_concat(jnp.arange(1, 5)), roll_existing(jnp.arange(1, 5)))
np.testing.assert_array_equal(roll_gather(jnp.arange(1, 5)), roll_existing(jnp.arange(1, 5)))

mesh = jax.sharding.Mesh(jax.devices(), 'i')
spec = jax.sharding.PartitionSpec('i')
sharding = jax.sharding.NamedSharding(mesh, spec)
x = jax.device_put(jnp.arange(16), sharding)

print(jax.jit(roll_existing).lower(x).compile().as_text())
print("-" * 80 + "\n")
print(jax.jit(roll_concat).lower(x).compile().as_text())
print("-" * 80 + "\n")
print(jax.jit(roll_gather).lower(x).compile().as_text())
```
<details>

```
HloModule jit_roll_existing, entry_computation_layout={(s32[4]{0})->s32[16]{0}}, allow_spmd_sharding_propagation_to_output={true}

%add (x: s32[], y: s32[]) -> s32[] {
  %x = s32[] parameter(0)
  %y = s32[] parameter(1)
  ROOT %add.5 = s32[] add(s32[] %x, s32[] %y)
}

%add.1 (x.1: s32[], y.1: s32[]) -> s32[] {
  %x.1 = s32[] parameter(0)
  %y.1 = s32[] parameter(1)
  ROOT %add.10 = s32[] add(s32[] %x.1, s32[] %y.1)
}

%fused_computation (param_0.1: u32[], param_1.2: u32[], param_2.5: s32[32]) -> s32[32] {
  %constant.62 = s32[] constant(0), metadata={op_name="jit(roll_existing)/jit(main)/jit(_roll)/jit(_roll)/concatenate[dimension=0]" source_file="<ipython-input-26-e19688a73297>" source_line=9}
  %broadcast.2 = s32[32]{0} broadcast(s32[] %constant.62), dimensions={}
  %param_2.5 = s32[32]{0} parameter(2)
  %param_0.1 = u32[] parameter(0)
  %convert.2 = s32[] convert(u32[] %param_0.1), metadata={op_name="jit(roll_existing)/jit(main)/jit(_roll)/jit(_roll)/concatenate[dimension=0]" source_file="<ipython-input-26-e19688a73297>" source_line=9}
  %constant.61 = s32[] constant(8), metadata={op_name="jit(roll_existing)/jit(main)/jit(_roll)/jit(_roll)/concatenate[dimension=0]" source_file="<ipython-input-26-e19688a73297>" source_line=9}
  %multiply.7 = s32[] multiply(s32[] %convert.2, s32[] %constant.61), metadata={op_name="jit(roll_existing)/jit(main)/jit(_roll)/jit(_roll)/concatenate[dimension=0]" source_file="<ipython-input-26-e19688a73297>" source_line=9}
  %dynamic-slice.8 = s32[8]{0} dynamic-slice(s32[32]{0} %param_2.5, s32[] %multiply.7), dynamic_slice_sizes={8}, metadata={op_name="jit(roll_existing)/jit(main)/jit(_roll)/jit(_roll)/concatenate[dimension=0]" source_file="<ipython-input-26-e19688a73297>" source_line=9}
  %param_1.2 = u32[] parameter(1)
  %constant.60 = u32[] constant(4)
  %multiply.6 = u32[] multiply(u32[] %param_1.2, u32[] %constant.60)
  %add.11 = u32[] add(u32[] %multiply.6, u32[] %param_0.1)
  %constant.59 = u32[] constant(8)
  %multiply.5 = u32[] multiply(u32[] %add.11, u32[] %constant.59)
  ROOT %dynamic-update-slice.3 = s32[32]{0} dynamic-update-slice(s32[32]{0} %broadcast.2, s32[8]{0} %dynamic-slice.8, u32[] %multiply.5)
}

%fused_computation.1 (param_0.3: s32[4], param_1.6: u32[]) -> s32[32] {
  %constant.65 = s32[] constant(0), metadata={op_name="jit(roll_existing)/jit(main)/jit(_roll)/jit(_roll)/concatenate[dimension=0]" source_file="<ipython-input-26-e19688a73297>" source_line=9}
  %broadcast.3 = s32[32]{0} broadcast(s32[] %constant.65), dimensions={}
  %param_0.3 = s32[4]{0} parameter(0)
  %param_1.6 = u32[] parameter(1)
  %convert.3 = s32[] convert(u32[] %param_1.6), metadata={op_name="jit(roll_existing)/jit(main)/jit(_roll)/jit(_roll)/concatenate[dimension=0]" source_file="<ipython-input-26-e19688a73297>" source_line=9}
  %constant.64 = s32[] constant(4), metadata={op_name="jit(roll_existing)/jit(main)/jit(_roll)/jit(_roll)/concatenate[dimension=0]" source_file="<ipython-input-26-e19688a73297>" source_line=9}
  %multiply.8 = s32[] multiply(s32[] %convert.3, s32[] %constant.64), metadata={op_name="jit(roll_existing)/jit(main)/jit(_roll)/jit(_roll)/concatenate[dimension=0]" source_file="<ipython-input-26-e19688a73297>" source_line=9}
  %dynamic-update-slice.5 = s32[32]{0} dynamic-update-slice(s32[32]{0} %broadcast.3, s32[4]{0} %param_0.3, s32[] %multiply.8), metadata={op_name="jit(roll_existing)/jit(main)/jit(_roll)/jit(_roll)/concatenate[dimension=0]" source_file="<ipython-input-26-e19688a73297>" source_line=9}
  %constant.63 = s32[] constant(16), metadata={op_name="jit(roll_existing)/jit(main)/jit(_roll)/jit(_roll)/concatenate[dimension=0]" source_file="<ipython-input-26-e19688a73297>" source_line=9}
  %add.12 = s32[] add(s32[] %multiply.8, s32[] %constant.63), metadata={op_name="jit(roll_existing)/jit(main)/jit(_roll)/jit(_roll)/concatenate[dimension=0]" source_file="<ipython-input-26-e19688a73297>" source_line=9}
  ROOT %dynamic-update-slice.4 = s32[32]{0} dynamic-update-slice(s32[32]{0} %dynamic-update-slice.5, s32[4]{0} %param_0.3, s32[] %add.12), metadata={op_name="jit(roll_existing)/jit(main)/jit(_roll)/jit(_roll)/concatenate[dimension=0]" source_file="<ipython-input-26-e19688a73297>" source_line=9}
}

ENTRY %main.41_spmd (param: s32[4]) -> s32[16] {
  %partition-id = u32[] partition-id(), metadata={op_name="jit(roll_existing)/jit(main)/jit(_roll)/jit(_roll)/concatenate[dimension=0]" source_file="<ipython-input-26-e19688a73297>" source_line=9}
  %replica-id = u32[] replica-id()
  %param = s32[4]{0} parameter(0), sharding={devices=[4]0,1,2,3}
  %fusion.1 = s32[32]{0} fusion(s32[4]{0} %param, u32[] %partition-id), kind=kLoop, calls=%fused_computation.1, metadata={op_name="jit(roll_existing)/jit(main)/jit(_roll)/jit(_roll)/concatenate[dimension=0]" source_file="<ipython-input-26-e19688a73297>" source_line=9}
  %all-reduce = s32[32]{0} all-reduce(s32[32]{0} %fusion.1), channel_id=1, replica_groups={{0}}, to_apply=%add, metadata={op_name="jit(roll_existing)/jit(main)/jit(_roll)/jit(_roll)/concatenate[dimension=0]" source_file="<ipython-input-26-e19688a73297>" source_line=9}
  %fusion = s32[32]{0} fusion(u32[] %partition-id, u32[] %replica-id, s32[32]{0} %all-reduce), kind=kLoop, calls=%fused_computation
  %all-reduce.1 = s32[32]{0} all-reduce(s32[32]{0} %fusion), channel_id=2, replica_groups={{0,1,2,3}}, use_global_device_ids=true, to_apply=%add.1
  ROOT %slice = s32[16]{0} slice(s32[32]{0} %all-reduce.1), slice={[15:31]}, metadata={op_name="jit(roll_existing)/jit(main)/jit(_roll)/jit(_roll)/dynamic_slice[slice_sizes=(16,)]" source_file="<ipython-input-26-e19688a73297>" source_line=9}
}


--------------------------------------------------------------------------------

HloModule jit_roll_concat, entry_computation_layout={(s32[4]{0})->s32[4]{0}}, allow_spmd_sharding_propagation_to_output={true}

%fused_computation (param_0: s32[1], param_1.1: s32[4]) -> s32[4] {
  %param_0 = s32[1]{0} parameter(0)
  %param_1.1 = s32[4]{0} parameter(1)
  %slice.4 = s32[3]{0} slice(s32[4]{0} %param_1.1), slice={[0:3]}, metadata={op_name="jit(roll_concat)/jit(main)/concatenate[dimension=0]" source_file="<ipython-input-26-e19688a73297>" source_line=12}
  ROOT %concatenate.1 = s32[4]{0} concatenate(s32[1]{0} %param_0, s32[3]{0} %slice.4), dimensions={0}, metadata={op_name="jit(roll_concat)/jit(main)/concatenate[dimension=0]" source_file="<ipython-input-26-e19688a73297>" source_line=12}
}

ENTRY %main.5_spmd (param: s32[4]) -> s32[4] {
  %param = s32[4]{0} parameter(0), sharding={devices=[4]0,1,2,3}
  %slice = s32[1]{0} slice(s32[4]{0} %param), slice={[3:4]}, metadata={op_name="jit(roll_concat)/jit(main)/concatenate[dimension=0]" source_file="<ipython-input-26-e19688a73297>" source_line=12}
  %collective-permute = s32[1]{0} collective-permute(s32[1]{0} %slice), channel_id=1, source_target_pairs={{0,1},{1,2},{2,3},{3,0}}, metadata={op_name="jit(roll_concat)/jit(main)/concatenate[dimension=0]" source_file="<ipython-input-26-e19688a73297>" source_line=12}
  ROOT %fusion = s32[4]{0} fusion(s32[1]{0} %collective-permute, s32[4]{0} %param), kind=kLoop, calls=%fused_computation, metadata={op_name="jit(roll_concat)/jit(main)/concatenate[dimension=0]" source_file="<ipython-input-26-e19688a73297>" source_line=12}
}


--------------------------------------------------------------------------------

HloModule jit_roll_gather, entry_computation_layout={(s32[4]{0})->s32[16]{0}}, allow_spmd_sharding_propagation_to_output={true}

%add.1 (x.1: s32[], y.1: s32[]) -> s32[] {
  %x.1 = s32[] parameter(0)
  %y.1 = s32[] parameter(1)
  ROOT %add = s32[] add(s32[] %x.1, s32[] %y.1)
}

%fused_computation (param_0.1: s32[4], param_1.4: u32[]) -> s32[16] {
  %constant.14 = s32[16,1]{1,0} constant({...})
  %param_1.4 = u32[] parameter(1)
  %convert.1 = s32[] convert(u32[] %param_1.4), metadata={op_name="jit(roll_gather)/jit(main)/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(), collapsed_slice_dims=(0,), start_index_map=(0,)) slice_sizes=(1,) unique_indices=True indices_are_sorted=False mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="<ipython-input-26-e19688a73297>" source_line=16}
  %constant.13 = s32[] constant(4), metadata={op_name="jit(roll_gather)/jit(main)/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(), collapsed_slice_dims=(0,), start_index_map=(0,)) slice_sizes=(1,) unique_indices=True indices_are_sorted=False mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="<ipython-input-26-e19688a73297>" source_line=16}
  %multiply.2 = s32[] multiply(s32[] %convert.1, s32[] %constant.13), metadata={op_name="jit(roll_gather)/jit(main)/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(), collapsed_slice_dims=(0,), start_index_map=(0,)) slice_sizes=(1,) unique_indices=True indices_are_sorted=False mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="<ipython-input-26-e19688a73297>" source_line=16}
  %broadcast.9 = s32[16,1]{1,0} broadcast(s32[] %multiply.2), dimensions={}, metadata={op_name="jit(roll_gather)/jit(main)/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(), collapsed_slice_dims=(0,), start_index_map=(0,)) slice_sizes=(1,) unique_indices=True indices_are_sorted=False mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="<ipython-input-26-e19688a73297>" source_line=16}
  %compare.3 = pred[16,1]{1,0} compare(s32[16,1]{1,0} %constant.14, s32[16,1]{1,0} %broadcast.9), direction=LT, metadata={op_name="jit(roll_gather)/jit(main)/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(), collapsed_slice_dims=(0,), start_index_map=(0,)) slice_sizes=(1,) unique_indices=True indices_are_sorted=False mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="<ipython-input-26-e19688a73297>" source_line=16}
  %constant.12 = s32[] constant(3), metadata={op_name="jit(roll_gather)/jit(main)/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(), collapsed_slice_dims=(0,), start_index_map=(0,)) slice_sizes=(1,) unique_indices=True indices_are_sorted=False mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="<ipython-input-26-e19688a73297>" source_line=16}
  %add.2 = s32[] add(s32[] %multiply.2, s32[] %constant.12), metadata={op_name="jit(roll_gather)/jit(main)/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(), collapsed_slice_dims=(0,), start_index_map=(0,)) slice_sizes=(1,) unique_indices=True indices_are_sorted=False mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="<ipython-input-26-e19688a73297>" source_line=16}
  %broadcast.8 = s32[16,1]{1,0} broadcast(s32[] %add.2), dimensions={}, metadata={op_name="jit(roll_gather)/jit(main)/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(), collapsed_slice_dims=(0,), start_index_map=(0,)) slice_sizes=(1,) unique_indices=True indices_are_sorted=False mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="<ipython-input-26-e19688a73297>" source_line=16}
  %compare.2 = pred[16,1]{1,0} compare(s32[16,1]{1,0} %constant.14, s32[16,1]{1,0} %broadcast.8), direction=GT, metadata={op_name="jit(roll_gather)/jit(main)/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(), collapsed_slice_dims=(0,), start_index_map=(0,)) slice_sizes=(1,) unique_indices=True indices_are_sorted=False mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="<ipython-input-26-e19688a73297>" source_line=16}
  %or.2 = pred[16,1]{1,0} or(pred[16,1]{1,0} %compare.3, pred[16,1]{1,0} %compare.2), metadata={op_name="jit(roll_gather)/jit(main)/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(), collapsed_slice_dims=(0,), start_index_map=(0,)) slice_sizes=(1,) unique_indices=True indices_are_sorted=False mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="<ipython-input-26-e19688a73297>" source_line=16}
  %reshape.9 = pred[16]{0} reshape(pred[16,1]{1,0} %or.2), metadata={op_name="jit(roll_gather)/jit(main)/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(), collapsed_slice_dims=(0,), start_index_map=(0,)) slice_sizes=(1,) unique_indices=True indices_are_sorted=False mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="<ipython-input-26-e19688a73297>" source_line=16}
  %constant.11 = s32[] constant(0), metadata={op_name="jit(roll_gather)/jit(main)/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(), collapsed_slice_dims=(0,), start_index_map=(0,)) slice_sizes=(1,) unique_indices=True indices_are_sorted=False mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="<ipython-input-26-e19688a73297>" source_line=16}
  %broadcast.7 = s32[16]{0} broadcast(s32[] %constant.11), dimensions={}, metadata={op_name="jit(roll_gather)/jit(main)/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(), collapsed_slice_dims=(0,), start_index_map=(0,)) slice_sizes=(1,) unique_indices=True indices_are_sorted=False mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="<ipython-input-26-e19688a73297>" source_line=16}
  %param_0.1 = s32[4]{0} parameter(0)
  %clamp.2 = s32[16,1]{1,0} clamp(s32[16,1]{1,0} %broadcast.9, s32[16,1]{1,0} %constant.14, s32[16,1]{1,0} %broadcast.8), metadata={op_name="jit(roll_gather)/jit(main)/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(), collapsed_slice_dims=(0,), start_index_map=(0,)) slice_sizes=(1,) unique_indices=True indices_are_sorted=False mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="<ipython-input-26-e19688a73297>" source_line=16}
  %subtract.1 = s32[16,1]{1,0} subtract(s32[16,1]{1,0} %clamp.2, s32[16,1]{1,0} %broadcast.9), metadata={op_name="jit(roll_gather)/jit(main)/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(), collapsed_slice_dims=(0,), start_index_map=(0,)) slice_sizes=(1,) unique_indices=True indices_are_sorted=False mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="<ipython-input-26-e19688a73297>" source_line=16}
  %gather.1 = s32[16]{0} gather(s32[4]{0} %param_0.1, s32[16,1]{1,0} %subtract.1), offset_dims={}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=1, slice_sizes={1}, metadata={op_name="jit(roll_gather)/jit(main)/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(), collapsed_slice_dims=(0,), start_index_map=(0,)) slice_sizes=(1,) unique_indices=True indices_are_sorted=False mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="<ipython-input-26-e19688a73297>" source_line=16}
  ROOT %select.1 = s32[16]{0} select(pred[16]{0} %reshape.9, s32[16]{0} %broadcast.7, s32[16]{0} %gather.1), metadata={op_name="jit(roll_gather)/jit(main)/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(), collapsed_slice_dims=(0,), start_index_map=(0,)) slice_sizes=(1,) unique_indices=True indices_are_sorted=False mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="<ipython-input-26-e19688a73297>" source_line=16}
}

ENTRY %main.4_spmd (param: s32[4]) -> s32[16] {
  %param = s32[4]{0} parameter(0), sharding={devices=[4]0,1,2,3}
  %partition-id = u32[] partition-id(), metadata={op_name="jit(roll_gather)/jit(main)/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(), collapsed_slice_dims=(0,), start_index_map=(0,)) slice_sizes=(1,) unique_indices=True indices_are_sorted=False mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="<ipython-input-26-e19688a73297>" source_line=16}
  %fusion = s32[16]{0} fusion(s32[4]{0} %param, u32[] %partition-id), kind=kLoop, calls=%fused_computation, metadata={op_name="jit(roll_gather)/jit(main)/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(), collapsed_slice_dims=(0,), start_index_map=(0,)) slice_sizes=(1,) unique_indices=True indices_are_sorted=False mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="<ipython-input-26-e19688a73297>" source_line=16}
  ROOT %all-reduce = s32[16]{0} all-reduce(s32[16]{0} %fusion), channel_id=1, replica_groups={{0}}, to_apply=%add.1, metadata={op_name="jit(roll_gather)/jit(main)/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(), collapsed_slice_dims=(0,), start_index_map=(0,)) slice_sizes=(1,) unique_indices=True indices_are_sorted=False mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="<ipython-input-26-e19688a73297>" source_line=16}
}
```

</details>
Do you have thoughts on how to implement this in an SPMD-friendly way?
The `roll_concat` implementation in my example snippet (basically the same as your initial suggestion) does generate SPMD friendly code -- it uses collective permute under the hood.
#16217 should address this.

Resolving patch:
diff --git a/jax/_src/numpy/lax_numpy.py b/jax/_src/numpy/lax_numpy.py
--- a/jax/_src/numpy/lax_numpy.py
+++ b/jax/_src/numpy/lax_numpy.py
@@ -3712,35 +3712,50 @@ def argpartition(a: ArrayLike, kth: int, axis: int = -1) -> Array:
 
 
 @partial(jit, static_argnums=(2,))
-def _roll(a, shift, axis):
-  a_shape = shape(a)
-  if axis is None:
-    return lax.reshape(_roll(ravel(a), shift, axis=0), a_shape)
-  shift = asarray(shift)
-  a_ndim = len(a_shape)
-  axis = np.asarray(axis)
-  b_shape = lax.broadcast_shapes(shift.shape, axis.shape, (1,))
+def _roll_dynamic(a: Array, shift: Array, axis: Sequence[int]) -> Array:
+  b_shape = lax.broadcast_shapes(shift.shape, np.shape(axis))
   if len(b_shape) != 1:
     msg = "'shift' and 'axis' arguments to roll must be scalars or 1D arrays"
     raise ValueError(msg)
 
   for x, i in zip(broadcast_to(shift, b_shape),
                   np.broadcast_to(axis, b_shape)):
-    i = _canonicalize_axis(i, a_ndim)
-    a_shape_i = array(a_shape[i], dtype=np.int32)
+    a_shape_i = array(a.shape[i], dtype=np.int32)
     x = ufuncs.remainder(lax.convert_element_type(x, np.int32),
-                  lax.max(a_shape_i, np.int32(1)))
-    a = lax.concatenate((a, a), i)
-    a = lax.dynamic_slice_in_dim(a, a_shape_i - x, a_shape[i], axis=i)
+                         lax.max(a_shape_i, np.int32(1)))
+    a_concat = lax.concatenate((a, a), i)
+    a = lax.dynamic_slice_in_dim(a_concat, a_shape_i - x, a.shape[i], axis=i)
   return a
 
+@partial(jit, static_argnums=(1, 2))
+def _roll_static(a: Array, shift: Sequence[int], axis: Sequence[int]) -> Array:
+  for ax, s in zip(*np.broadcast_arrays(axis, shift)):
+    if a.shape[ax] == 0:
+      continue
+    i = (-s) % a.shape[ax]
+    a = lax.concatenate([lax.slice_in_dim(a, i, a.shape[ax], axis=ax),
+                         lax.slice_in_dim(a, 0, i, axis=ax)],
+                        dimension=ax)
+  return a
 
 @util._wraps(np.roll)
-def roll(a, shift, axis: Optional[Union[int, Sequence[int]]] = None):
-  util.check_arraylike("roll", a,)
-  if isinstance(axis, list):
-    axis = tuple(axis)
-  return _roll(a, shift, axis)
+def roll(a: ArrayLike, shift: Union[ArrayLike, Sequence[int]],
+         axis: Optional[Union[int, Sequence[int]]] = None) -> Array:
+  util.check_arraylike("roll", a)
+  arr = asarray(a)
+  if axis is None:
+    return roll(arr.ravel(), shift, 0).reshape(arr.shape)
+  axis = _ensure_index_tuple(axis)
+  axis = tuple(_canonicalize_axis(ax, arr.ndim) for ax in axis)
+  if not core.is_constant_shape(arr.shape):
+    # TODO(necula): support static roll for polymorphic shapes.
+    return _roll_dynamic(arr, asarray(shift), axis)
+  try:
+    shift = _ensure_index_tuple(shift)
+  except TypeError:
+    return _roll_dynamic(arr, asarray(shift), axis)
+  else:
+    return _roll_static(arr, shift, axis)
 
 
 @util._wraps(np.rollaxis, lax_description=_ARRAY_VIEW_DOC)


Test patch: