Is communication w/ sharding really this easy?

[maedoc]

When knocking my head against the wall with pmap over a simple sparse matrix example, I started reading the docs on sharding and expected it to be harder. My question is, have I misunderstood what's going on? am I going to hit some snag on the communication in the background? Here's a MWE

import os, numpy, scipy.sparse
os.environ['XLA_FLAGS'] = '--xla_force_host_platform_device_count=8'
import jax, jax.numpy as np
from jax.experimental import mesh_utils
from jax.sharding import PositionalSharding

# sparse matrix vector multiply as a learning example
A = scipy.sparse.rand(512, 512, density=0.03, format='coo', dtype='f')
x = np.ones(A.shape[0])
y1 = A @ x
y2 = np.zeros(A.shape[0]).at[A.row].add(A.data * x[A.col])
numpy.testing.assert_allclose(y2, y1)

# setup sharding & sharded arrays
shard = PositionalSharding(mesh_utils.create_device_mesh((8, )))
y3_ = np.zeros(A.shape[0])
r_, c_, d_, x_, y_ = [
    jax.device_put(np.array(_), shard)
    for _ in (A.row, A.col, A.data, x, y3_)]
print('x_ shard')
jax.debug.visualize_array_sharding(x_)

# sharded to common array
y3 = np.zeros(A.shape[0]).at[r_].add(d_ * x_[c_])
print('y3 shard, allclose? ', np.allclose(y3, y1, 1e-6, 1e-6))
jax.debug.visualize_array_sharding(y3)

# sharded to sharded array
y4 = jax.device_put(np.zeros(A.shape[0]), shard)
y4 = y4.at[r_].add(x_[c_] * d_)
print('y4 shard, allclose? ', np.allclose(y4, y1, 1e-6, 1e-6))
jax.debug.visualize_array_sharding(y4)

As far as I can tell, this actually works for parallelizing a sparse matrix-vector multiply, and there's no communication written out explicitly. If it's really this easy, it would perhaps share some examples, since I don't recall seeing anything int he docs about semantics of x.at[idx].add in sharding context and it completely surprised me that it worked (for both y3 and y4 above).

edit to add the output of the above script

x_ shard
┌───────┬───────┬───────┬───────┬───────┬───────┬───────┬───────┐
│ CPU 0 │ CPU 1 │ CPU 2 │ CPU 3 │ CPU 4 │ CPU 5 │ CPU 6 │ CPU 7 │
└───────┴───────┴───────┴───────┴───────┴───────┴───────┴───────┘
y3 shard, allclose?  True
┌───────────────────┐
│CPU 0,1,2,3,4,5,6,7│
└───────────────────┘
y4 shard, allclose?  True
┌───────┬───────┬───────┬───────┬───────┬───────┬───────┬───────┐
│ CPU 0 │ CPU 1 │ CPU 2 │ CPU 3 │ CPU 4 │ CPU 5 │ CPU 6 │ CPU 7 │
└───────┴───────┴───────┴───────┴───────┴───────┴───────┴───────┘

[jcandane]

@maedoc Did you end up finding the answer? Is it possible this is just a lucky result. I would imagine: Sharding the Sparse-Index-Arrays, might lead to situations in which their index does not correspond to the Sharded-Dense-Array (e.g. I shard the index array, and the 4th component contains an index which is specified on the 6th dense sharded part).
For dense arrays, I think it is that easy.

[maedoc]

@jcandane I think the result is that yes, the placement of computation really does just follow the placement of data, as promised. The downside, for my use cases, was that the performance is not nearly as good as I expected, at least on CPU, e.g. a basic CSR SpMV routine in C multithreaded with OpenMP is more performant. Other experiments with Jax yielded the conclusion that XLA (the LLVM backend that Jax targets) seems pretty good at GPU, but for high performance CPU, one can look elsewhere: in one example of Monte Carlo simulations, the NumPy code is indeed 4x slower than Jax, but Jax is 4-8x slower than a C implementation which uses SIMD effectively. Even tho the data layout is the same in the Jax and C code, so that SIMD can be used, XLA seems generate mainly scalar code, which, for CPU, is a big performance loss.