pytorch: how to apply function over all cells of 4-d tensor

Question

I'm trying to apply a function over a 4-D tensor (I think about it as a 2-D matrix with a 2-D matrix in each cell) with the following dimensions: [N x N x N x N]. The apply function returns [1 x N] tensor, so after the apply function I'm expecting a tensor of the following dimensions: [N x N x 1 x N].

Example: let's define [4 x 4 x 4 x 4] tensor:

tensor_4d = torch.tensor([[[[1., 0., 0., 0.],
  [0., 0., 0., 0.],
  [0., 0., 0., 0.],
  [0., 0., 0., 0.]],
 [[1., 0., 0., 0.],
  [1., 0., 0., 0.],
  [0., 0., 0., 0.],
  [0., 0., 0., 0.]],
 [[1., 0., 0., 0.],
  [0., 0., 0., 0.],
  [0., 0., 0., 0.],
  [0., 0., 0., 0.]],
 [[1., 0., 0., 0.],
  [0., 0., 0., 0.],
  [0., 0., 0., 0.],
  [0., 0., 0., 0.]]],
[[[0., 0., 0., 0.],
  [0., 1., 0., 0.],
  [0., 0., 0., 0.],
  [0., 0., 0., 0.]],
 [[0., 0., 0., 0.],
  [0., 1., 0., 0.],
  [0., 0., 0., 0.],
  [0., 0., 0., 0.]],
 [[0., 0., 0., 0.],
  [0., 1., 0., 0.],
  [0., 0., 0., 0.],
  [0., 0., 0., 0.]],
 [[0., 0., 0., 0.],
  [0., 1., 0., 0.],
  [0., 0., 0., 0.],
  [0., 0., 0., 0.]]],
[[[0., 0., 1., 0.],
  [0., 0., 0., 0.],
  [0., 0., 1., 0.],
  [0., 0., 0., 0.]],
 [[0., 0., 0., 0.],
  [0., 0., 1., 0.],
  [0., 0., 1., 0.],
  [0., 0., 0., 0.]],
 [[0., 0., 0., 0.],
  [0., 0., 0., 0.],
  [0., 0., 1., 0.],
  [0., 0., 0., 0.]],
 [[0., 0., 0., 0.],
  [0., 0., 0., 0.],
  [0., 0., 1., 0.],
  [0., 0., 0., 0.]]],
[[[0., 0., 1., 0.],
  [0., 0., 0., 0.],
  [0., 0., 0., 1.],
  [0., 0., 0., 1.]],
 [[0., 0., 0., 0.],
  [0., 0., 1., 0.],
  [0., 0., 0., 1.],
  [0., 0., 0., 1.]],
 [[0., 0., 0., 0.],
  [0., 0., 0., 0.],
  [0., 0., 0., 1.],
  [0., 0., 0., 1.]],
 [[0., 0., 0., 0.],
  [0., 0., 0., 0.],
  [0., 0., 0., 0.],
  [0., 0., 0., 1.]]]], dtype=torch.float64)

lets look on tensor_4d at [3][0]:

tensor_4d[3][0]
tensor([[0., 0., 1., 0.],
        [0., 0., 0., 0.],
        [0., 0., 0., 1.],
        [0., 0., 0., 1.]], dtype=torch.float64)

this my apply function :

def apply_function(tensor_2d):
    eigenvalues, eigenvectors = torch.eig(input=tensor_2d, eigenvectors=True)
    return eigenvector[:, 2]

and this the result of the apply function:

apply_function(tensor_4d[3][0])
tensor([-1.0000e+00,  0.0000e+00, 4.0083e-292,  0.0000e+00],
       dtype=torch.float64)

so the apply_function works for each cell. Next, I'm trying to use the apply_function with the whole matrix, and expecting each cell will contain the result of activating 'apply_function' for this cell. but, when using the apply function I'm getting the following error:

apply_function(tensor_4d)

Traceback (most recent call last):
  File "C:\Program Files\JetBrains\PyCharm Community Edition 2019.2.3\helpers\pydev\_pydevd_bundle\pydevd_exec2.py", line 3, in Exec
    exec(exp, global_vars, local_vars)
  File "<input>", line 1, in <module>
  File "C:/Users/LiavB/PycharmProjects/RBC/Components/RBC_torch.py", line 41, in apply_function
    eigenvalues, eigenvectors = torch.eig(input=tensor_2d, eigenvectors=True)
RuntimeError: invalid argument 1: A should be 2 dimensional at ..\aten\src\TH/generic/THTensorLapack.cpp:206

Answer 1

Let's try:

new_shape=(-1,)+tensor_4d.shape[2:]

out = (torch.stack([apply_function(t) for t in tensor_4d.view(new_shape)], axis=-1)
      .reshape(new_shape)
)