Adj_jac_apply

The magic happens in something like this: https://github.com/stan-dev/math/blob/develop/stan/math/rev/mat/fun/simplex_constrain.hpp#L95

You pass the functor as a template argument to adj_jac_apply and the autodiff arguments as regular arguments. It hooks everything in with autodiff and returns you some vars.

Sure, but there is in addition always a function defined for the double only case. At least this is what is the case for the softmax… and the function you refer to always returns a var type of thing. So this function cannot just return an Eigen::VectorXd thing.

Oh oh sorry, yes, you’re right. You still need to define the prim version (mostly cause that’ll also handle fvars and fvar<var>s and whatever else)

Oh no… I just realized that, if you have a Jacobian already, then using adj_jac_apply to as a way to inject the Jacobian into the ODE solver is gonna be really inefficient.

The issue is you have the Jacobian, supply it to adj_jac_apply, and then the jacobian function in the ODE solver will reconstruct the jacobian one matrix vector product at a time (where only one element of these vectors is non-zero).

precomputed_gradients will be a better option for what you’re doing.

but that adj_jac_apply facility looks far more convenient to code up as I have to deal not that much with the var stuff. Moreover, if there are a ton of zeros I can take advantage of that.

Let’s see how this will come along.

+1. As soon as we get it tested, it’s totally the way to calculate directional derivatives. I don’t know how well that’ll plug into the ODE solvers, but it’ll be faster even for systems with as many outputs as inputs (or even slightly fewer given the lower constant factors).

One issue is that we haven’t optimized any of our forward-mode derivatives. So the gain is in theory and should hold for simple things, but I wouldn’t trust it for complicated matrix stuff (though even then, plugging in the custom types should be reasonable).

You certainly won’t need it for programming, but the union card helps for jobs at universities and research labs.

A 10% speedup is “nice”. 4x is at least “great”, if not “fantastic”. It helps to log scale the y axis for performance reports so we can read off the speedup and see how it depends on N.

Yes. The design is set up to be mutable after the functor call, then immutable until the adjoint call, so you can save intermediate values for use in the Jacobian. You just need to use variables on our autodiff stack.

Yes, indeed. Having zeros in double-based matrix arithmetic’s no problem as they don’t incur overhead. At the same time, the adjoint-Jacobian-apply thing isn’t really set up for sparsity of the whole adjoint vector, as it assumes each adjoint needs to be incremented.

The ideal is to let the implementation of the adjoint-Jacobian product be customizable so that if a function has a sparse Jacobian then its autodiff implementation can take advantage of that, right?

Yurp, that’s right.

Right—it can deal with sparse Jacobians. It just can’t deal with sparseness in the adjoint-Jacobian product result. All those entries get filled in and incremented.

@bbbales2 saw this from another branch. If you want to use google benchmark for this stuff I have a branch of perf-math here with your original example.

You can make the original test with

make adj_jac_apply_add.o && make adj_jac_apply_add
./adj_jac_apply_add
# run multiple to get summary stats
./adj_jac_apply_add --benchmark_repetitions=15 --benchmark_report_aggregates_only=true

Oh cool. Yeah I wonder if there’s a way forwarding might speed all this up a little. @andrjohns was just looking at adj_jac_apply.

1 Like

(At this point in the meeting you aren’t gonna miss much if you bail – just the roundtable is useful to have a vague idea of what everyone else is working on)

(Also thanks for coming! It’s just a weekly reminder of what everyone is working on and what they’re thinking!)

1 Like

Hey @stevebronder, can I borrow your forwarding experience when you’ve got a minute? I’m working on adding forwarding (and generalising eigen inputs) to adj_jac_apply and want to check that I haven’t missed anything (still a new area for me).

I’ve got the changes up in this branch, would you mind having a quick review? I’ve definitely over-forwarded in some areas (scalars), but may have missed other things (or done them less than efficiently). Thanks!

Sure lmk when it’s ready!

Lol those two messages above were meant for a PM. Good work me.

1 Like

Lol oh gosh sorry I totally misread this I will look at it tmrw

Haha you’re all good! No rush!