Given a 2D tensor A
of size m x r
, and a list of steps steps = [m_1, ..., m_k]
that partition m
(ie, sum(steps)=m
).
Define max_m_i = \max_i m_i
and s_i = sum(steps[:i])
.
I need to create a 3D tensor of size k x max_m_i x r
by padding every slice of A
w.r.t. steps, ie, A[s_i:s_i+steps[i], :]
of size m_{i+1} x r
pad to the size max_m_i x r
.
Example,
import torch
m = 6; r = 2
A = torch.arange(0, 12).reshape(6, 2)
print("A = ", A)
steps = [3, 2, 1]
print(f"{steps = }")
max_m_i = max(steps)
pad_A = []
s_i = 0
for i in range(len(steps)):
pad_A += [torch.cat((A[s_i:s_i + steps[i]], torch.zeros((max_m_i - steps[i], r))), axis=0)]
s_i += steps[i]
pad_A = torch.stack(pad_A)
print("pad_A = ", pad_A)
output
A = tensor([[ 0, 1],
[ 2, 3],
[ 4, 5],
[ 6, 7],
[ 8, 9],
[10, 11]])
steps = [3, 2, 1]
pad_A = tensor([[[ 0., 1.],
[ 2., 3.],
[ 4., 5.]],
[[ 6., 7.],
[ 8., 9.],
[ 0., 0.]],
[[10., 11.],
[ 0., 0.],
[ 0., 0.]]])