Adds lb/ub/lub_constrain specializations #2373
Conversation
stan/math/rev/fun/lb_constrain.hpp
Outdated
| auto is_inf_lb = to_arena( | ||
| (arena_lb.val().array() != NEGATIVE_INFTY).template cast<double>()); | ||
| auto precomp_x_exp = to_arena((arena_x.val().array() * is_inf_lb).exp()); | ||
| arena_t<ret_type> ret = (is_inf_lb).select( | ||
| precomp_x_exp + arena_lb.val().array(), arena_x.val().array()); | ||
| reverse_pass_callback( | ||
| [arena_x, arena_lb, ret, is_inf_lb, precomp_x_exp]() mutable { | ||
| arena_x.adj().array() += ret.adj().array() * precomp_x_exp; | ||
| arena_lb.adj().array() += ret.adj().array() * is_inf_lb; | ||
| }); |
There was a problem hiding this comment.
This is sort of a dirty trick, but since exp(0) = 1 We can cast all of these checks to double and then in the forward pass multiply x by the check array. When we make ret we iterate over is_inf_lb and when it's true it uses the values we want and when false just returns arena_x at that value.
But the nice thing is that we can then have no checks in the reverse pass and get nice SIMD instructions since we are just doing an elementwise multiplication
arena_x.adj().array() += ret.adj().array() * precomp_x_exp;
There was a problem hiding this comment.
PRs introducing such dirty tricks should be accompanied by benchmarks proving they are actually faster than the alternative implementation.
There was a problem hiding this comment.
This doesn't seem like a dirty trick. Maybe just a trick lol. I'm happy with it.
There was a problem hiding this comment.
Actually this ended being best case the same speed and worst case a bit slower so I reverted it to just do the checks
Yes, it should be. Do you even need holders for constrain/free functions? I think probably not. |
I don't like the expressions there since lp is an output in this code. Even if it weren't segfaulting, that would mean an accidental double eval of the expression would make a math bug rather than just a performance bug (updating lp twice).
I thought x was always either a scalar or a matrix here. |
|
@SteveBronder I made a comparison branch where I removed a lot of stuff: feature/lb_constrain-matvar...feature/lb_constrain-matvar_shrunk I deleted a bunch of stuff and I think this will work okay too. I don't think we need to use identity constrain in here. I think we only need reverse mode specializations when both types are matrices, otherwise prim can handle the varmat autodiff. Ofc. this would have more temporaries, but the implementation is way simpler. Can also go verbose for the performance -- it's just a metric ton of code. |
|
idt those will work because we need to check each cell individually for infinity. Doing it the way you wrote it above you end up just checking if all of them are positive infinity and if not then do the constrain. |
The signatures are: (and then the ones for lp) I'm pretty sure the code I wrote covers that. Is there a signature I'm missing here? |
|
I think that could work but I already wrote out the complicated bits lol so I'd rather just use the specializations since they would be more efficient |
|
Makes sense. Lemme revieeeeew |
stan/math/prim/fun/lb_constrain.hpp
Outdated
| if (unlikely(lb == NEGATIVE_INFTY)) { | ||
| return identity_constrain(x, lb); | ||
| } else { | ||
| // check_less("lb_constrain", "lb", value_of(x), value_of(lb)); |
stan/math/prim/fun/lb_constrain.hpp
Outdated
| if (lb == NEGATIVE_INFTY) { | ||
| return identity_constrain(x, lb); | ||
| } else { | ||
| // check_less("lb_constrain", "lb", value_of(x), value_of(lb)); |
stan/math/prim/fun/lb_constrain.hpp
Outdated
| [](const auto& x, const auto& lb) { | ||
| return x.unaryExpr([lb](auto&& x) { return lb_constrain(x, lb); }); | ||
| }, | ||
| std::forward<T>(x), std::forward<L>(lb)); |
There was a problem hiding this comment.
[optional] I'd rather just evaluate expressions here because I'm pretty sure we're always assigning to l-values somewhere in generated C++. I think it does cost a malloc or something to make a holder? @t4c1 can confirm if they're free or not.
There was a problem hiding this comment.
Yeah good point
There was a problem hiding this comment.
If a holder is constructed from lvalues only it is virtually free. We need a malloc for each rvalue. But let me ask again: Do we even need holders here?
|
|
||
| stan::test::expect_ad(f1, x, lb); | ||
| // stan::test::expect_ad(f1, x, lb); |
| using stan::math::lb_constrain; | ||
| using stan::math::promote_scalar_t; | ||
| Eigen::MatrixXd A(2, 2); | ||
| // Doesn't work for values near zero fail for fvar hessian? |
…4.1 (tags/RELEASE_600/final)
…4.1 (tags/RELEASE_600/final)
…4.1 (tags/RELEASE_600/final)
|
@rok-cesnovar hey sorry but can we do the reverse of the last constrain pr up in Stan? We just need to revert the changes to the lb/ub constrain checks |
|
Actually @rok-cesnovar we are going to merge all the changes for the constraints into this PR with sub PRs then once all those are here we can revert the Stan PR |
Adds ub_constrain for varmat
…4.1 (tags/RELEASE_600/final)
…4.1 (tags/RELEASE_600/final)
…h into feature/lb_constrain-matvar
…4.1 (tags/RELEASE_600/final)
Add specializations for lub_constrain for var<matrix>
SteveBronder
changed the title
|
I put up #3012 in the Stan repo. @rok-cesnovar how should we do the double mere thing? |
…4.1 (tags/RELEASE_600/final)
Jenkins Console Log Machine informationProductName: Mac OS X ProductVersion: 10.11.6 BuildVersion: 15G22010CPU: G++: Clang: |
|
Nevermind, problem, don't merge yet |
Jenkins Console Log Machine informationProductName: Mac OS X ProductVersion: 10.11.6 BuildVersion: 15G22010CPU: G++: Clang: |
|
No sorry this is ready for merge I made a false alarm. I thought this was valid Stan code parameters {
real xx[N];
real<lower=xx,upper=1> theta[N, N];
}but it's not so we r good |
Jenkins Console Log Machine informationProductName: Mac OS X ProductVersion: 10.11.6 BuildVersion: 15G22010CPU: G++: Clang: |
Summary
This PR adds the specializations for
lub_constrainfor matrices and scalars with all their combinations of inputs. The branch I was working on to make lb/ub/lub constrain work for matvar was getting pretty big so I'm breaking it up into 3 seperate ones.Tests
Tests were updated for
matvarand to the calculations when computing the gradient of both lp and the input.Side Effects
There's a few weirds I can't completely sort out.
@t4c1 Is
make_holderable to deal with plain references? I was trying to use it for the functions that need to calculate lp likeBut I was getting seg faults when trying.
@andrjohns is it possible to use the apply meta stuff to have this work for
std::vectorsof containers? I couldn't figure out how to have it work whenever I have anxlikestd::vector<Eigen::MatrixXd>and anlbofdoubleRelease notes
Replace this text with a short note on what will change if this pull request is merged in which case this will be included in the release notes.
Checklist
Math issue How to add static matrix? #1805
Copyright holder: Steve Bronder
The copyright holder is typically you or your assignee, such as a university or company. By submitting this pull request, the copyright holder is agreeing to the license the submitted work under the following licenses:
- Code: BSD 3-clause (https://opensource.org/licenses/BSD-3-Clause)
- Documentation: CC-BY 4.0 (https://creativecommons.org/licenses/by/4.0/)
the basic tests are passing
./runTests.py test/unit)make test-headers)make test-math-dependencies)make doxygen)make cpplint)the code is written in idiomatic C++ and changes are documented in the doxygen
the new changes are tested