def _allgather(x, dim, size, axis_name):
shape = list(x.shape)
shape.insert(dim, size)
out = lax.full(shape, lax._const(x, 0))
out = lax.dynamic_update_index_in_dim(out, x, axis_index(axis_name), dim)
return psum(out, axis_name)
def _empty_array(sz, aval):
if aval is core.abstract_unit:
return core.unit
else:
return lax.full((sz,) + aval.shape, 0, aval.dtype)
def _expand(dim, size, axis_name, x):
shape = list(x.shape)
shape.insert(dim, size)
out = lax.full(shape, lax._const(x, 0))
return lax.dynamic_update_index_in_dim(out, x, axis_index(axis_name), dim)
def _empty_arrays(aval):
assert isinstance(aval, core.AbstractValue)
if isinstance(aval, core.AbstractTuple):
return core.pack(map(_empty_arrays, aval))
else:
return lax.full(aval.shape, 0, aval.dtype)
