Reshape original X from LSTM with predictions
I have a tensor with my LSTM inputs X (in PyTorch) as well as the matching predictions Y_hat
I want to add the Y_hat as a column to the original X.
The problem is that LSTM a sliding window with seq_length. In case seq. length is 3 and I have 6 variables in X, and 2 variables in Y_hat, I have something like this:
First entry of Tensor:
X1: 
[1 2 3 4 5 6]
[5 4 6 7 8 9]
[3 6 8 7 4 8]

Y_hat1
[0 1]

X2:
[5 4 6 7 8 9] (repeated from X1)
[3 6 8 7 4 8] (repeated from X1)
[4 8 7 9 8 4]
Y_hat1
[1 1]

and so on.
Is there an easy pythonesk command to reshape the X with Y so I can quickly get:
[3 6 8 7 4 8 0 1]
[4 8 7 9 8 4 1 1] 


Having defined X1, X2, Y_hat1, and Y_hat2:
X1 = torch.tensor([[1, 2, 3, 4, 5, 6],
                   [5, 4, 6, 7, 8, 9],
                   [3, 6, 8, 7, 4, 8]])
Y_hat1 = torch.tensor([0,1])

X2 = torch.tensor([[5, 4, 6, 7, 8, 9],
                   [3, 6, 8, 7, 4, 8],
                   [4, 8, 7, 9, 8, 4]])
Y_hat2 = torch.tensor([1,1])

As well as stacks of inputs (X) and targets (Y):
X = torch.stack([X1, X2])
Y = torch.stack([Y_hat1, Y_hat2])

You can the desired operation using concatenation and stacking operators:
&gt;&gt;&gt; X = torch.stack([X1, X2])
&gt;&gt;&gt; Y = torch.stack([Y_hat1, Y_hat2])
&gt;&gt;&gt; torch.stack(tuple(torch.cat((x[-1], y)) for x, y in zip(X, Y)))

Which expanded corresponds to:
&gt;&gt;&gt; torch.stack((
        torch.cat((X1[-1], Y_hat1)),
        torch.cat((X2[-1], Y_hat2))))

In a more vectorized fashion you can do:
&gt;&gt;&gt; torch.hstack((X[:,-1], Y))

