Как объединить назначение фрагментов, назначение маски и трансляцию в PyTorch? ⇐ Python

[Anonymous]

Чтобы быть более конкретным, мне интересно, как одновременно назначить тензор по срезу и по маске в разных измерениях в PyTorch.
Вот небольшой пример того, что я хочу сделать:
С тензорами и масками ниже:

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

x = torch.zeros(2, 3, 4, 6)
mask = torch.tensor([[ True, True, False], [True, False, True]])
y = torch.rand(2, 3, 1, 3)

Я хочу добиться чего-то вроде

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

x[mask, :, :3] = y[mask]

• В измерениях 0 и 1 только фрагменты 4x6/1x3 в x/

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

  y

  тому, чей соответствующий элемент в маске имеет значение True, разрешено присваивать.
• Я надеюсь, что в измерении 2 тензор из 1 строки в y может транслироваться на все 8 строк x,
• В измерении 3 только первым трем элементам x присваиваются 3 -тензор элемента из y.

Однако в приведенном выше коде была обнаружена следующая ошибка:

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

RuntimeError: shape mismatch: value tensor of shape [4, 1, 3] cannot be broadcast to indexing result of shape [4, 3, 6]

Похоже, что PyTorch вместо этого выполнил индексацию по [маске] и проигнорировал индексацию :3.
Я' я тоже пробовал

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

x[mask][:, :, :3] = y[mask]

Ошибок не произошло, но назначение по-прежнему не выполнено.
Я знаю, что могу назначать по срезу и по маске шаг за шагом, но надеюсь избежать любые промежуточные тензоры, если это возможно.
Тензоры в нейронных сетях могут быть очень большими, поэтому выполнение комплексного присваивания может занять меньше времени и меньше памяти.

Подробнее здесь: https://stackoverflow.com/questions/791 ... in-pytorch