Как реализовать логику решения задачи с тензорами без циклов (срезами)

Постановка задачи

Имеется mapper - вспомогательная структура для заполнения массива modificated_arr значениями из condition1 и condition2. Маппер это двумерный массив (одинаковой размерности с modifacted_arr) заполненный списками из двух элементов, где первый указывает на индекс массива из conditions, а второй - какой именно массив (condition1 или condition2)/ Например элемент [3, 0] - означает, что по индексам 0, 1 массива modificated_arr - нужно установить значение 4 элемента из condition1.Решение задачи простое - проитерироваться по mapper-у, считать данные и заполнить массив modicated_arr этими значениями. Проблема заключается в том, что надо это сделать для тензоров без итераций, то есть назначать значения сразу срезам массива modificated_arr. Типа такого modificated_arr[:,:] = mappper[slice1,slice2,slice3]

Каждая из переменных ниже представляет собой тензор torch, но думаю решение numpy массивов не сильно будет отличаться. Поэтому добавлен тег numpy

mapper = torch.tensor([
          [[0,0],[3, 0], [4, 1]],
          [[2,1],[4, 0], [3, 1]]
         ])

condition1 = torch.tensor([0, 1, 0, 0, 1])
condition2 = torch.tensor([1, 0, 1, 1, 0])

modificated_arr = torch.tensor([
                   [0, 0, 0],
                   [0, 0, 0]
                  ])

Решение с циклами:

for row_idx, row in enumerate(mapper):
    for col_idx, col in enumerate(row):

        cond_idx, cond = col

        if cond == 0:
            modificated_arr[row_idx, col_idx] = condition1[cond_idx]

        else:
            modificated_arr[row_idx, col_idx] = condition2[cond_idx]

Дополнения к объяснению:

Первые два измерения mapper-а соответствуют измерениям modificated_arr. То есть mapper[1,2] должен будет изменить элемент массива modificated_arr[1,2] следующим образом. mapper[1,2] это массив [3,1]. Первый элемент двумерного массива означает номер индекса condition1 или condition2, второй элемент - 1 отвечает какой из condition - ов брать, он соответствует condtion2 то есть нужно будет изменить modificated_arr[1,2] = cond2[3].


Ответы (1 шт):

Автор решения: Хачатур Саркисян

Решение можно осуществить с булевыми масками.

modificated_arr[mapper[:, :, 1] == 0] += condition1[mapper[:, :, 0]][mapper[:, :, 1] == 0]
modificated_arr[mapper[:, :, 1] == 1] += condition2[mapper[:, :, 0]][mapper[:, :, 1] == 1]

Получаем ожидаемый результат:

modificated_arr
Out[6]: 
tensor([[0, 0, 0],
        [1, 1, 1]])
→ Ссылка