We are getting there… yes I think that will do what you wrote. We need those two operations: ReduceF for parameters -> function output/gradients and CombineF for merging things into the AD stack.
From looking more closely on the
Eigen::Tensor library I think we should use it. It allocates memory at a single chunk in column-major format. Hence, we can convert data into this format and then just slice it up for the different nodes. On each node we then can hopefully use a reduction operation which uses their framework, see here:
I am not yet sure if we can use their reduce thing, because we have a dynamically sized output reduce operation as we do not know in advance how many elements are being returned from some function call. I need to check that (in case we cannot use Eigen’s reduction thing, then we just write our own).
So on each node we generate in step 4 another Eigen Tensor which is put together on the root node and the step 6 can be represented as a reduction operation as well, I think.
Let me try to get the Eigen reductions working and then I hopefully can sketch the next version.
Ahh… if you don’t mind, I would try to avoid templates where we can. For example, if we go with
Eigen::Tensor then we will only ever have to deal with those objects of type double (only the integer data must stay as nested vector thing).
Just one more thought: We could make our life a bit easier in implementing a map_rect_lpdf version first. This version has the advantage that we know that only a single value per function call is returned. Returning only the log-lik value is anyway what I would recommend anyone when using MPI as this cuts down communication cost. You have any thought on this? Pro: Easier for us to implement + anyway fast; Con: not as flexible and possible even more cumbersome to program for the user.
BTW…while I don’t think we have to burden ourselfes with this as a requirement, we should be able to use this design for the ODE integrators as well (write function outputs and their gradients to an
Eigen::Tensor and putting that onto the AD stack).