torch.block_diag

torch.block_diag(*tensors) [исходный код]

Создаёт блочно-диагональную матрицу из переданных тензоров.

Параметры:

*tensors – Один или несколько тензоров размерности 0, 1 или 2.

Возвращает:

Двумерный тензор, содержащий все входные тензоры, расположенные по порядку так, что их верхние левые и нижние правые углы являются соседними по диагонали. Все остальные элементы равны 0.

Тип возвращаемого значения:

Тензор

Пример:

>>> import torch
>>> A = torch.tensor([[0, 1], [1, 0]])
>>> B = torch.tensor([[3, 4, 5], [6, 7, 8]])
>>> C = torch.tensor(7)
>>> D = torch.tensor([1, 2, 3])
>>> E = torch.tensor([[4], [5], [6]])
>>> torch.block_diag(A, B, C, D, E)
tensor([[0, 1, 0, 0, 0, 0, 0, 0, 0, 0],
        [1, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        [0, 0, 3, 4, 5, 0, 0, 0, 0, 0],
        [0, 0, 6, 7, 8, 0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0, 7, 0, 0, 0, 0],
        [0, 0, 0, 0, 0, 0, 1, 2, 3, 0],
        [0, 0, 0, 0, 0, 0, 0, 0, 0, 4],
        [0, 0, 0, 0, 0, 0, 0, 0, 0, 5],
        [0, 0, 0, 0, 0, 0, 0, 0, 0, 6]])