Stack tensors with different shapes in Eigen across dim 1

Question

I have a series of tensors of shape [B, n, C], where B and C are of constant size, and n varies between tensors (say, n in [1, 5]). I would like to concatenate all of these tensors into a single one of shape [B, k, C], where k is the sum of all values of n. In numpy, this would be:

import numpy as np

batch_size = 2
num_channels = 4

values = []

for i in range(0, 5):
  tensor = np.ones([batch_size, i+1, num_channels]) * i
  values.append(tensor)

print(np.concatenate(values, axis=1))

With the result

[[[0. 0. 0. 0.]
  [1. 1. 1. 1.]
  [1. 1. 1. 1.]
  [2. 2. 2. 2.]
  [2. 2. 2. 2.]
  [2. 2. 2. 2.]
  [3. 3. 3. 3.]
  [3. 3. 3. 3.]
  [3. 3. 3. 3.]
  [3. 3. 3. 3.]
  [4. 4. 4. 4.]
  [4. 4. 4. 4.]
  [4. 4. 4. 4.]
  [4. 4. 4. 4.]
  [4. 4. 4. 4.]]

 [[0. 0. 0. 0.]
  [1. 1. 1. 1.]
  [1. 1. 1. 1.]
  [2. 2. 2. 2.]
  [2. 2. 2. 2.]
  [2. 2. 2. 2.]
  [3. 3. 3. 3.]
  [3. 3. 3. 3.]
  [3. 3. 3. 3.]
  [3. 3. 3. 3.]
  [4. 4. 4. 4.]
  [4. 4. 4. 4.]
  [4. 4. 4. 4.]
  [4. 4. 4. 4.]
  [4. 4. 4. 4.]]]

How can I achieve this in Eigen?

Answer 1

I found a solution using slicing:

#include <Eigen/Core>
#include <Eigen/Dense>
#include <iostream>
#include <vector>

#include "eigen3/unsupported/Eigen/CXX11/Tensor"

using IndexType = EIGEN_DEFAULT_DENSE_INDEX_TYPE;
using TensorType3D = Eigen::Tensor<float, 3, Eigen::RowMajor, IndexType>;

int main() {
  // Initialization:
  // Tensors of size [B, i+1, C], with i in [0, 4]. Each of this
  // has i as its value.
  TensorType3D output(2, 15, 4);
  std::vector<TensorType3D> values;
  for (int i = 0; i < 5; ++i) {
    TensorType3D input(2, i + 1, 4);
    for (int j = 0; j < 2; ++j) {
      for (int m = 0; m < i + 1 + 1; ++m) {
        for (int k = 0; k < 4; ++k) {
          input(j, m, k) = i;
        }
      }
    }
    values.push_back(input);
  }
  
  // Assign values via slicing.
  int previous = 0;
  for (int i = 0; i < values.size(); ++i) {
    const int size = values[i].dimension(1);
    output.slice(Eigen::array<Eigen::Index, 3>{0, previous, 0},
                 Eigen::array<Eigen::Index, 3>{2, size, 4}) = values[i];
    previous += size;
  }

  std::cout << output << "\n";

  return 0;
}

The output is equivalent to Python's:

0 0 0 0
1 1 1 1
1 1 1 1
2 2 2 2
2 2 2 2
2 2 2 2
3 3 3 3
3 3 3 3
3 3 3 3
3 3 3 3
4 4 4 4
4 4 4 4
4 4 4 4
4 4 4 4
4 4 4 4

0 0 0 0
1 1 1 1
1 1 1 1
2 2 2 2
2 2 2 2
2 2 2 2
3 3 3 3
3 3 3 3
3 3 3 3
3 3 3 3
4 4 4 4
4 4 4 4
4 4 4 4
4 4 4 4
4 4 4 4

Answer 2

You can use .concatenate:

#include <unsupported/Eigen/CXX11/Tensor>
#include <iostream>

int main(){
    Eigen::Tensor<int,3> A(2,5,4);
    Eigen::Tensor<int,3> B(2,7,4);
    Eigen::Tensor<int,3> C = A.concatenate(B, 1); // Concatenate along axis 1 of A
    std::cout << C.dimensions();
}

which prints [2, 12, 4].

Compiler explorer