This builds on a previous conversation about custom C++ implementations of multivariate functions: Vector-valued functions with manual gradients in external C++
I have successfully got the attached small example working of a (trivial effectless) function with specialised custom adjoints for both ‘vector of vars’ and ‘var vector’ (matvar/varmat) input types. This compiles and runs fine. There are probably various sophistications missing so that this setup will fail in certain edge cases. Any suggestions for improving this are welcome, of course, but my key question is slightly different.
In the full model that this is ultimately for, the adjoint calculation is rather complex, with subtle indexing gymnastics. I would therefore prefer not having to deal with keeping two (matvar and varmat) versions in different locations, since they only differ in how the results of those calculations are fed into the vector/matrix of adjoints. There are obviously various ways of avoiding this (including a mid-file #include), but one approach that seems like it should be possible is to have a combined varmat/matvar version of the function.
My attempt to do so hit a roadblock. It is commented out in the .hpp file. Basically, I cannot figure out how to structure the type traits to trigger on both varvec and vecvar input. The trivial way of combining the individual requires I already have would seem to be some kind of OR, but the requires aren’t boolean. Intuitively, it seems as if return_type should work, but that also fails.
Another question in this context is whether it is ok to return a (more efficient?) varvec regardless of whether the input was varvec or vecvar?
Any help with this would be much appreciated, as I’m facing the frustrating situation with the full model where I only got a matvar version with custom adjoints to work, but it runs more slowly than just letting Stan do its AD magic. So I’m really hoping a proper varmat version of custom adjoints will actually give a speed boost.
Thanks!
multigrad_test.hpp (2.4 KB)
multigrad_test.stan (232 Bytes)
multigrad_test.R (285 Bytes)
This is generally how we do it in the Math library, yes. We use require_rev_matrix(/_vector) which accepts both types. This file also shows both ways to do a disjunction, either with stan::math::disjunction or require_any_t, using the is_XYZ templates rather than the require_XYZ templates.
No, this will not be handled well.
Note that unless you intend on submitting this code to the Math library or running your own modified version of the compiler, it will currently not generate SoA calls to the code anyway (c.f. [BUG] SoA request is ignored with external c++ · Issue #1379 · stan-dev/stanc3 · GitHub)
Thanks, @WardBrian , that’s really helpful!
Good to know about require_rev_matrix as a convenient solution.
I thought returning a varmat would be allowed because it says somewhere that a function accepting an autodiff type should also return "an" autodiff type, which sounded like it need not necessarily be the same.
The last point you mention is disappointing. As I mentioned, with my full model the custom AoS version with known analytic adjoints for some reason runs more slowly than a naive version providing no explicit adjoints and leaving it up to the Stan’s AD. I was hoping the SoA approach would fix that. Is that to some extent expected when analytic adjoints get complicated, or does it sound like there might be something other than AoS inefficiency going on and I should post that full model in a separate thread?
I have no intention of making the code fit for official inclusion in the Math library at this point. Having said that, the function is fairly generic, in that calculates the cell probabilities when you put a grid on a bivariate normal. If you think that’s of general interest and potentially worth including in the library, I’d be happy to contribute what I can, but would probably require some support.
It certainly wouldn’t hurt to see the full thing! In general I’d say that if your hand-coded derivative is slower than just letting autodiff go through the primal, that’s probably an indication of something that can be improved
The full code would also let us think about if it would make sense for inclusion in the math library. It sounds like you’ve already gotten quite close to what would be a solid first PR
Thanks a lot!
Ok, so the full model is attached here. But having moved to require_rev_matrix as input types resulted in it no longer compiling. My guess is the problem is related to the return type of the function overall. I’m not sure how to specify it in terms of the input type so that the template actually defines two different specializations for matvar and varmat, given that you said forcing both to the same output type is problematic. Anyhow, perhaps I’m mistaken and the problem is something else entirely.
Note that I’m making use of the external approxcdf package. If this is a challenge in terms of reproducability, the dependency could be avoided for testing purposes by replacing the call to norm_cdf_2d_vfast with a proxy multiplication of the marginals (this alternative has been commented out). The numbers would be wrong, of course, as this would effectively ignore rho, but right now I’m focused on getting the type gymnastics right.
I call the attached files from within R as
model <- cmdstanr::cmdstan_model(
'src/stan/bvn_cpp_test.stan',
stanc_options = list("allow-undefined" = TRUE),
user_header = normalizePath('src/stan/bvn_all_probs.hpp'),
force_recompile = TRUE
)
Any suggestions would be much appreciated!
bvn_cpp_test.hpp (5.6 KB)
bvn_cpp_test.stan (469 Bytes)
@bifouba I was able to get your code compiling with only two small changes:
- When you calculate res, you’re turning it into an
Eigen::MatrixXd and then calling to_arena, but this breaks any relationship between the types of the input and the type of the output. Instead I updated the code to read stan::arena_t<T1> res_arena = bvn_all_probs(stan::math::value_of(ect), stan::math::value_of(lct), stan::math::value_of(rho), pstream__);. Having a type dependency on the input T1 ensures res_arena is an autodiff type of compatible memory layout
- You’re returning
res, not res_arena
I also combined the two matrix template types into 1, since it’s rarely a good idea to mix them. I’ve attached my version which is compiling (note I did have to comment out the approxcdf stuff – thanks for the tip)
bvn_probs.hpp (5.7 KB)
Thank you so much, @WardBrian !
I can confirm that this does compile on my end. So I was pretty close after all. Especially since I was actually correctly returning res_arena originally, but must have switched it to plain res during a phase of blind trial and error… I understand the point about needing to maintain the type relationship. With hindsight, that’s one point where previously using make_callback_var meant this was taken care of elsewhere I at the time that I looked at your fft example code I hadn’t realised that plain_type_t leaves the “var-ness” intact. Anyhow, I think I understand what you did and why that works, so thanks again! Btw. I should clarify that calling this the “full model” was misleading. I meant it’s the full C++ function. But the Stan model exercising it is just a minimal test set up, no the real thing.
One snag is that although this version compiles, when I actually run the model, it throws a runtime error complaining that
Assertion failed: (v == T(Value)), function variable_if_dynamic, file XprHelper.h, line 133.
Since this is about dynamic (types?), let me just double-check that this set-up should in theory work even if the first two inputs and the output are the same T1 type, but have distinct dynamic sizes?
I would try a version that has the extra template parameter again. That’s not an assertion I’ve seen before, but it seems like something where Eigen is using an expression template type for one but not the other, so forcing them to be the same was an error on my part
Thanks again! I had to both give each function argument its own template parameter again but in addition had to assign res_arena the type stan::arena_t<stan::return_var_matrix_t<Eigen::MatrixXd, T0, T1>>, but now it seems to work!
It does indeed work, but unfortunately, the hand-coded derivatives appears to yield no discernible performance gain compared to just letting autodiff do its thing. I suppose if you had noticed anything obviously grossly inefficient in the gradient implementation, you would have mentioned it already, @WardBrian , so perhaps that’s just a reflection of the complexity of the analytic gradient and the quality of autodiff! Although given that they are more or less exactly equally fast even on large matrices, I do wonder how I know that the custom gradients actually do get used and if there is a way of checking that?
Nothing stood out to me code wise, but that isn’t to say there isn’t a more efficient way of calculating the same gradient analytically that I would have missed. It is quite a complicated one!
The way you’ve got it coded now it would have to be using the custom reverse pass, because the primal isn’t templated enough to possibly accept vars. If you’re only ever calling it in autodiff contexts, you could rename the primal to something not matching the name you use in the Stan model to make this 100% clear, but it really shouldn’t be necessary.
Out of curiosity, are you just comparing times using total sample time from the model, or are you using something like profile blocks
In general, when benchmarking actual model variants we try to identify both the raw computation cost of a single iteration in purely computational and the cost per effective sample (which may be different because of lower acceptance rate, autocorrelation etc.). But since in this case both variants of the function (with and without gradients) return identical function values, there shouldn’t be any fit-related differences and I just used total run time as a proxy for computational efficiency. But we did identify this function as a target for custom implementation based on proper profiling.