Любой стиль NumPy / факел, чтобы установить значение с учетом индекса ndarray и флаг ndarray?

Question

Я работаю над PyTorch, и в настоящее время я столкнулся с проблемой, для которой я понятия не имею, как решить ее в стиле факела / numpy. Например, предположим, что у меня есть три тензора PyTorch

import torch
import numpy as np

indices = torch.from_numpy(np.array([[2, 1, 3, 0], [1, 0, 3, 2]]))
flags = torch.from_numpy(np.array([[False, False, False, True], [False, False, True, True]]))
tensor = torch.from_numpy(np.array([[2.8, 0.5, 1.2, 0.9], [3.1, 2.8, 1.3, 2.5]]))

Здесь flags - это тезис логического флага, показывающий, какие элементы в indices должны быть извлечены. Учитывая извлеченные индексы, я хочу установить соответствующие элементы в tensor в указанное const (скажем, 1e-30). Исходя из приведенного выше примера, я хочу

>>> sub_indices = indices.op1(flags)
>>> sub_indices
tensor([[0], [3, 2]])
>>> tensor.op2(sub_indices, 1e-30)
>>> tensor
tensor([[1e-30, 0.5, 1.2, 0.9], [3.1, 2.8, 1e-30, 1e-30]])

Может ли кто-нибудь помочь дать решение? Я использую понимание списка, но я думаю, что это немного уродливо. Я попытался indices[flags], но он возвращает только 1d-массив [0, 3, 2], поэтому применение этого приведет к изменению всех строк в тех же столбцах 0, 2, 3

Некоторые дополнительные замечания:

• Количество истинных значений для каждой строки в flags не может быть определено
• Каждая строка indices обязательно является перестановкой последовательности 0 ... N - 1

Ниже приводится примерная версия кода для удобства вставки копий. Я сомневаюсь, что это можно сделать чисто тупым способом

import numpy as np

indices = np.array([[2, 1, 3, 0], [1, 0, 3, 2]])
flags = np.array([[False, False, False, True], [False, False, True, True]])
tensor = np.array([[2.8, 0.5, 1.2, 0.9], [3.1, 2.8, 1.3, 2.5]])

Answer 1

Pytorch версия решения @ soloice. В pytorch torch.gather используется вместо torch.take.

indices = torch.tensor([[2, 1, 3, 0], [1, 0, 3, 2]])
flags = torch.tensor([[False, False, False, True], [False, False, True, True]])
tensor = torch.tensor([[2.8, 0.5, 1.2, 0.9], [3.1, 2.8, 1.3, 2.5]])

indices_sorted = indices.argsort(axis=1)
mask = torch.gather(flags, 1, indices_sorted).float()
result = tensor * (1 - mask) + 1e-30 * mask

Answer 2

Вы можете отсортировать flags в соответствии с indices, чтобы создать mask, а затем использовать mask в качестве мультиплексора. Вот пример кода:

indices = np.array([[2, 1, 3, 0], [1, 0, 3, 2]])
flags = np.array([[False, False, False, True], [False, False, True, True]])
tensor = np.array([[2.8, 0.5, 1.2, 0.9], [3.1, 2.8, 1.3, 2.5]])

indices_sorted = indices.argsort(axis=1)
mask = np.take_along_axis(flags, indices_sorted, axis=1)
result = tensor * (1 - mask) + 1e-30 * mask

Я не совсем знаком с pytorch, но думаю, что собирать рваный тензор не очень хорошая идея. Хотя, даже в худшем случае, вы можете конвертировать в / из массивов numpy.