Problems raised when trying to pad or roll a sharded array

[zhiping0913]

When I am trying to pad a sharded array, under some modes, there will have errors.

import jax
jax.config.update('jax_num_cpu_devices', 2)
print(jax.local_device_count())
print(jax.devices())
import jax.numpy as jnp
from jax.sharding import PartitionSpec,NamedSharding
mesh = jax.make_mesh((2,), ('x', ))
sharding = NamedSharding(mesh, PartitionSpec('x', ))
print(sharding)

arr=jnp.arange(16)
arr_shard = jax.device_put(arr, sharding)
arr_shard_pad=jnp.pad(arr_shard, [(2, 2)],mode='reflect')

pad function works when mode is 'constant', but it will have error when mode is 'edge' or 'reflect'.

Traceback (most recent call last):
  File "/1.py", line 23, in <module>
    arr_shard_pad=jnp.pad(arr_shard, [(2, 2)],mode='reflect')
  File "/.conda/envs/mpi_python/lib/python3.14/site-packages/jax/_src/numpy/lax_numpy.py", line 4372, in pad
    return _pad(array, pad_width, mode, constant_values, stat_length, end_values, reflect_type)
  File "/.conda/envs/mpi_python/lib/python3.14/site-packages/jax/_src/numpy/lax_numpy.py", line 4198, in _pad
    return _pad_symmetric_or_reflect(array, pad_width, str(mode), reflect_type)
  File "/.conda/envs/mpi_python/lib/python3.14/site-packages/jax/_src/numpy/lax_numpy.py", line 4048, in _pad_symmetric_or_reflect
    array = build_padding(array, pad_width[i][0], before=True)
  File "/.conda/envs/mpi_python/lib/python3.14/site-packages/jax/_src/numpy/lax_numpy.py", line 3998, in build_padding
    edge = lax_slicing.slice_in_dim(array, 0, 1, axis=i)
jax._src.core.ShardingTypeError: slicing on sharded dims where out dim (1) is not divisible by mesh axes (2) with spec (x) is not implemented.

Similar error will be raised when I am trying to roll a sharded array.
When the shaft in roll is shift=-2, then this can work.

arr_shared_roll=jnp.roll(arr_shard, shift=-2, axis=0)

However, when roll is an odd number, then it failed.

arr_shared_roll=jnp.roll(arr_shard, shift=-1, axis=0)
Traceback (most recent call last):
  File "/1.py", line 24, in <module>
    arr_shared_roll=jnp.roll(arr_shard, shift=-1, axis=0)
  File "/.conda/envs/mpi_python/lib/python3.14/site-packages/jax/_src/numpy/lax_numpy.py", line 8570, in roll
    return _roll_static(arr, shift, axis)
  File "/.conda/envs/mpi_python/lib/python3.14/site-packages/jax/_src/numpy/lax_numpy.py", line 8514, in _roll_static
    a = lax.concatenate([lax_slicing.slice_in_dim(a, i, a.shape[ax], axis=ax),
                         ~~~~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^
jax._src.core.ShardingTypeError: slicing on sharded dims where out dim (15) is not divisible by mesh axes (2) with spec (x) is not implemented.

I know I can ensure the pad and roll function work if I bring it back to a normal, single-device array.

arr_back=jax.device_get(arr_shard)
arr_pad=jnp.pad(arr_back, [(2, 2)],mode='reflect')
arr_roll=jnp.roll(arr_back, shift=-1, axis=0)

This will make the program less efficiency.
How to pad or roll a sharded array?

[jakevdp]

I think essentially this is failing when the operation results in implicit cross-device communication. This is something that can be a silent performance bottleneck in large workflows, and so the compiler errors when it encounters an object that requires it.

In the most general case, the easiest way to get around this would probably be to explicitly reshard the array such that it's replicated on every device, then perform the roll, and then reshard the output again to the desired output sharding. That's a potentially expensive operation in terms of data movement and storage, which is why the compiler doesn't do it silently. But given the data movement required for rolling a sharded array, I'm not sure if there's a better approach except, in specific cases where the sharding aligns exactly with the size of the roll (like your even-sized roll with 2 devices).

Maybe @yashk2810 has ideas?