stan-dev
[FR] Supporting Eigen::Tensor
Description
It would be a lot of work, but I think if we supported Eigen's Tensor types we could use those for much more efficient array structures.
In Stan, when we have something like
data {
matrix[N, M] Y[3, 4, 5];
}
What we actually have in the C++ is something like
std::vector<std::vector<std::vector<Eigen::MatrixXd>>> Y;
Under the hood this translates to a big structure of pointers pointing to other pointers. In pseudocode the above structure is almost something like
double**** Y;
This also makes it a bit trickier to give optimization advice because for matrices we are column major, but it's actually more optimal to access arrays as if we are going over row major constructed memory.
Eigen's Tensor module supports multi-dimensional objects like we currently have for arrays. For the example above
Eigen::Tensor<double, 5> t(N, M, 3, 4, 5);
This is nice because now our data structure is simply a double* with memory in column major order. Another benefit of this is with the new matrix type we can use our current scheme and place an Eigen Tensor into the var like
var_value<Eigen::Tensor<double, 5>>
There's a lot of things to sort out for this to work. The docs are here. Essentially anywhere that currently accepts an std::vector<> needs to accept an Eigen::Tensor. Some functions will need new specializations, but other functions that already accept either Eigen::Vector or std::vector<> types should require minimal changes. A good first step may be to get the distributions working with 1 dimensional tensors.
For the Stan repo we would need reading in and out tensors from a var_context as well as operating on them via the serializer and deserializer. For stanc we would need to change the current code that uses std::vector<> to generate an Eigen::Tensor<> instead.
I personally am not sure I'd have time to work on this, but for anyone that is interested I'd be happy to do reviews!
Current Version:
v4.1.0
Some other considerations are that we would need to use the plugin extension and add to_matrix(), to_vector(), and to_row_vector() so if we have a Stan structure like
matrix[N, M] Y[2, 3, 4];
Once we access the outer dimensions of the array part the stanc compiler need to know it is now working with a matrix. I'm putting a little godbolt below with a tensor setup if anyone wants to play with it
https://godbolt.org/
I really like this! Do we need to be concerned that Tensor is in the unsupported modules?
Ty! It's maintained by Google so idt so. It's only in unsupported because it requires c++11