extension of all_gather

[mganahl]

Hi,
I wanted propose a feature extension for all_gather. Our team is currently running test-calculations on TPU pods, and we use a custom broadcasting operation build in JAX that might also be useful for others.

We'd like to propose the addition of an optional source_indices argument to all_gather that would allow
gathering values from a particular source index to all indices within one element in axis_index_groups, i.e.

A = np.arange(4)
print(pmap(lambda x: jax.lax.all_gather(x, 
   axis_name='i', 
   axis_index_groups=[[0,2],[1,3]], 
   source_indices=[0,3])) #[0,3,0,3]

This is a minimal implementation that achieves this:

def mask(A, cond):
    do_not_mask = jnp.zeros_like(A, dtype=np.bool) + cond
    return jnp.where(do_not_mask, x=A, y=jnp.zeros_like(A))

 
def all_gather(x, axis_name, axis_index_groups=None, source_indices=None):
    if source_indices is None:
        return jax.lax.all_gather(x, axis_name=axis_name, 
                                  axis_index_groups=axis_index_groups)
    assert len(source_indices) == len(axis_index_groups)
    for n, source in enumerate(source_indices):
        assert source in axis_index_groups[n]
        
    axis_index = jax.lax.axis_index(axis_name)
    keep =jax.numpy.any(jnp.array([axis_index == s for s in source_indices]))
    masked =  mask(x, keep)
    return jax.lax.psum(masked,axis_name,axis_index_groups=axis_index_groups)
    
    
    

[apaszke]

So it's really more like a broadcast from a single device to all other devices in its axis index group? I don't think that fits very well with the semantics of all_gather, so it should at least be a separate function. But, perhaps more importantly, I think that we should start moving away from axis_index_groups, in favor of xmap. So I think that we might be interested in merging something similar to what you propose, but only if we can make sense of it in the world of multidimensional meshes and named axes. Do you think that makes sense?

[chaserileyroberts]

I believe once this is accepted into stablehlo, it should be exactly what you want. I plan on adding this to Jax under lax.pcollective_broadcast once this RFC is accepted.

openxla/stablehlo#1809