I'm trying to distribute large computation to multiple shards using jax. I have 2x2 tiles of shards. Each shard has 3x3 array (with distinct integers), e.g. in shard (0,0): x = jnp.array([[40, 54, 70], [14, 3, 94], [44, 54, 88]]) and in shard (1,0): x = jnp.array([[50, 46, 22], [ 4, 83, 21], [91, 32, 68]]) I wanted to use lax.gather() to obtain first column from the sibling shard so that shard (0,0) will get [50, 4, 91] and other shard will get [40, 14, 44]. And then I can use the column from another shard to update first column in the local shard to keep min values. Another possible approach is to use pmin, but I did not find working example that could be used for the purpose. My experimental code: first_column = lax.gather( operand=x, start_indices=jnp.array([0]), dimension_numbers=jax.lax.GatherDimensionNumbers( offset_dims=(0,), start_index_map=(1,), collapsed_slice_dims=()), slice_sizes=jnp.array([x.shape[0]])) returns error: TypeError: Shapes must be 1D sequences of concrete values of integer type, got On TFRT_CPU_0 at mesh coordinates (i, j,) = (0, 0): [3] Continue reading...
