Simplify tensor-matrix operation with numpy.einsum or einops

265 Views Asked by At

I have the following inputs:

  • T: a (H x W x C) tensor, or if needed (H x W x C x 1).
  • M: (C x C) matrix.

I need to compute a (H x W x C) tensor, in which each "slice" is the matrix product between M and the corresponding (1 x 1 x C) slice from T, like this:

result = np.zeros_like(T)
for row in range(H):
    for col in range(W):
        result[row, col, ...] = M @ T[row, col, ...]

enter image description here

Can this be done more efficiently with numpy, numpy.einsum or einops?

0

There are 0 best solutions below