Получение столбца от другого осколка с использованием jax.lax.gather () ⇐ Python

[Anonymous]

Я пытаюсь распространять большие вычисления на несколько осколков, используя JAX. У меня есть 2x2 плитки осколков. Каждый осколок имеет массив 3x3 (с различными целыми числами), например, в Shard (0,0): < /p>

Код: Выделить всё

x = jnp.array([[40, 54, 70], [14,  3, 94], [44, 54, 88]])
< /code>
и в Shard (1,0): < /p>
x = jnp.array([[50, 46, 22], [ 4, 83, 21], [91, 32, 68]])

Я хотел использовать lax.gather () для получения первого столбца от брата -осколка, чтобы Shard (0,0) получит [50, 4, 91], а другой Shard получит [40, 14, 44]. И затем я могу использовать столбец из другого осколка для обновления первого столбца в локальном Shard, чтобы сохранить значения MIN. < /P>
Другой возможный подход - это использование PMIN, но я не нашел рабочий пример, который можно было бы использовать для этой цели. < /P>
Мой экспериментальный код: < /p>
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]]))
< /code>
Возвращает ошибку: < /p>

typeerror: формы должны быть 1D -последовательностями конкретных значений целочисленного типа, получили на tfrt_cpu_0 в координатах сетки (i, j,) = (0, 0):
[3] < /p> /> Изменить:
Я думаю, что я нашел обходной путь с помощью all_gather (): < /p>
j_coord = lax.axis_index('j')
column = lax.cond(j_coord == 0, lambda _: x[:,-1], lambda _: x[:,0], operand=None)
column_gathered = lax.all_gather(column, axis_name="j", tiled=False)

column_gathered = jnp.minimum(column_gathered[0,:], column_gathered[1,:])

x = lax.cond(j_coord == 0,
lambda _: x.at[:,-1].set(column_gathered),
lambda _: x.at[:,0].set(column_gathered), operand=None)
< /code>
Однако я обеспокоен тем, что All_gather () отправляет два столбца вместо одного. Есть ли способ сделать расчет минимума среди осколков более эффективным, чем это?

Подробнее здесь: https://stackoverflow.com/questions/796 ... lax-gather