def testAxisIndex(self):
device_count = xla_bridge.device_count()
f = pmap(lambda x: x + pxla.axis_index('i'), 'i')
x = np.ones(device_count)
ans = f(x)
expected = 1 + onp.arange(device_count)
self.assertAllClose(ans, expected, check_dtypes=False)
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 _drop(x, dim, axis_name):
return lax.dynamic_index_in_dim(x, axis_index(axis_name), dim, False)
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)