Получение столбца от другого осколка с использованием 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):
< /p>
< /blocktote.>

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