Repository navigation
Matrix comprehensions (build_matrix function) #3282
Description
Activity
Sounds awesome. My first thought was isn't this close to "apply" in R?
I don't think the function in
applygets access to the indices, but broadly speaking it's a similar idea. It's also close to python's list comprehensions -- in fact, semantically, we basically wantbuild_vectorto be equivalent todef build_vector(f, N, *args): return [f(i, *args) for i in range(N)]
So I started toying with this today. I was looking at how
integrate_1dandlaplacehandle precalculating the gradients. One issue is that both of those functions always return a scalar where we want back a matrix. Since laplace andintegrate_1donly have one output we only need one pass of reverse mode to pre calculate the jacobian adjoint product info we need. But thisbuild_matrixneeds to return a matrix based on several several inputs. So I tried doing an adjoint method likecvodes_integrator_adjoint.hpp, but running benchmarks I'm finding that to be much slower.I have my version here that uses a reverse pass inside of the reverse pass to calculate the adjoints we need.
https://github.com/stan-dev/math/compare/feat/build_matrix?expand=1
The graph below shows the slowdown relative to a simple loop. So the simple loop is almost always faster.
I’m not necessarily surprised it’s slower than a simple loop in a microbenchmark, it would need to also be compared with the benefit of then using an SoA result in a further computation vs the loop being stuck in AoS
icic, if the goal is just to allow more promotion to SoA then I think we should get rid of all of the adjoint tricks I have and just have a simple for loop version. Then at the end I'll promote it to SoA. That will always be good.
tbh I think a thing more beneficial than this would be a for loop collapser optimization pass in the compiler. We can use the monotonic framework over for loops to detect if we can collapse simple for loops into vectorized instructions
vector[N] x; // ... vector[N] y; for (i in 1:N) { y[i] = exp(x[i]); }
into
vector[N] y = exp(x);
If we want to force promotion to SoA we could also do it through stan-dev/stanc3#1439 with an optimization pass that checks if a vector/matrix after a loop should be promoted to SoA (which would just make a new
{name}_soa_tmp__used in the rest of the program.@SteveBronder I talked to @avehtari about this in the language meeting this morning. He said he was unsurprised that there would not be any speedup if the number of parameters was the same as the size of the matrix.
The use case he is interested in is for something like a gp covariance matrix, which has a smaller number of free parameters going in than the NxM matrix coming out (e.g., the entire matrix could be computed from data, except for the diagonal elements which depend on some parameters)
It would be worth benchmarking some examples that align with this use case more, if possible. I've asked @avehtari to supply some
This would be a variadic function with signature:
matrix build_matrix(int I, int J, Fn f, T1 arg1,...TN argN)(name up for debate)where
Fnis a function(int, int, T1, ..., TN) => realIt returns a matrix
Msuch thatM[i,j] = f(i,j, arg1, ... argN)The real advantage to this function is that it can do local autodiff, which will improve memory usage/locality and allow us to have the returned matrix be a SoA matrix (cf #2344 (comment)).
Once implemented,
build_[row]_vectoris the obvious extension.