Get `sd`, `log_softmax`, and `log_sum_exp` fully `var<mat>` compatible (Issue #2101) #2169

bbbales2 · 2020-10-27T20:03:18Z

Summary

This makes sd, log_sum_exp, and log_softmax fully var<mat> compatible (and apply_vector_unary in the process).

Release notes

Updated sd, log_softmax, and log_sum_exp to work fully with var<mat>

Checklist

Math issue Make functions with custom autodiff var<mat> friendly #2101
Copyright holder: Columbia University

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
- unit tests pass (to run, use: ./runTests.py test/unit)
- header checks pass, (make test-headers)
- dependencies checks pass, (make test-math-dependencies)
- docs build, (make doxygen)
- code passes the built in C++ standards checks (make cpplint)
the code is written in idiomatic C++ and changes are documented in the doxygen
the new changes are tested

…4.1 (tags/RELEASE_600/final)

bbbales2

Questions & comments

stan/math/prim/functor/apply_vector_unary.hpp

bbbales2 · 2020-10-27T20:05:40Z

stan/math/prim/meta/is_container.hpp

@@ -19,7 +20,8 @@ namespace stan {
 */
 template <typename Container>
 using is_container = bool_constant<
-    math::disjunction<is_eigen<Container>, is_std_vector<Container>>::value>;
+  math::disjunction<is_eigen<Container>, is_std_vector<Container>,
+		    is_var_matrix<Container>>::value>;


A var<mat> should count as a container, right? This might break template logic elsewhere but we can fix it elsewhere.

stan/math/rev/fun/sd.hpp

test/unit/math/mix/fun/sd_test.cpp

SteveBronder · 2020-10-28T02:32:35Z

I'll take a look at this in the morning

bbbales2 · 2020-11-02T20:42:12Z

Ping

… `sd` (Issue #2101)

…4.1 (tags/RELEASE_600/final)

…into feature/sd_varmats

…4.1 (tags/RELEASE_600/final)

bbbales2

@SteveBronder this is ready to look at. Not 100% there but I have a couple questions. One we get those ironed out it will make it easy to finish mat<var> ing log_sum_exp and log_softmax.

stan/math/prim/fun/sd.hpp

test/unit/math/test_ad.hpp

SteveBronder

Couple quick comments

stan/math/rev/fun/sd.hpp

test/unit/math/mix/fun/sd_test.cpp

test/unit/math/test_ad.hpp

…mat> types (Issue #2101)

bbbales2 · 2020-11-19T20:45:18Z

stan/math/rev/fun/log_softmax.hpp

+    alpha.adj().noalias()
+      += res.adj_
+      - (res.adj_.sum() * res.val_.array().exp()).matrix();
+  });


@t4c1, @SteveBronder writing this code I wanted a make_callback_var to wrap make_callback_vari. The argument pass to the functor can still be the vari to avoid the pointer chasing.

I also want .adj() and val() on vari_value.

That all sound okay to add? (Edit: if so I'll just do it here)

I want make_callback_var to avoid type algebra. In this case the input and output types will match, in other cases they won't and I would end up writing:

using ret_type = decltype((theta.array() - log(theta.exp().sum())); return var_value<ret_type>(...);

…4.1 (tags/RELEASE_600/final)

bbbales2 · 2020-11-19T20:57:34Z

test/unit/math/test_ad.hpp

+          require_std_vector_vt<is_matrix, Type>* = nullptr>
+void check_return_type(const ReturnType& ret, const Type& x) {
+  if (ret.size() > 0 && x.size() > 0)
+    check_return_type(ret[0], x[0]);


I will not be offended if you want me to take a closer look at check_return_type.

I think the logic we want is:

If there are only var_value<double> s on the input, then there should only be var_value<double> s on the output

If there are var_value<not double> s on the input, then there should be no var_value<double> s on the output

I think this would need an extra template program var_value_t to extract the var_value from a generic input type and I am too lazy to write it today.

I think the logic is fine here unless I'm missing something. This is a specialization for std::vector<T> that just checks that the inner type for the input/output is correct. If that's what this does then I think it's fine

…t a matrix (Issue #2101)

SteveBronder

Just a couple little things to change and then this looks good.

SteveBronder · 2020-11-20T21:33:31Z

stan/math/rev/fun/log_softmax.hpp

+  return make_callback_vari(
+      (theta.array() - log(theta.exp().sum())).matrix(),
+      [alpha](const auto& res) mutable {
+        alpha.adj().noalias()
+            += res.adj_ - (res.adj_.sum() * res.val_.array().exp()).matrix();
+      });


[side note] It would be nice to have a make_callback_var so we could still use .adj() and .val() etc. Mostly a slice of life feature

I added this

SteveBronder · 2020-11-20T21:33:46Z

stan/math/rev/fun/log_softmax.hpp

+T log_softmax_impl(const T& alpha) {
+  check_nonzero_size("log_softmax", "alpha", alpha);
+
+  const auto& theta = to_ref(alpha.val().array() - alpha.val().maxCoeff());


Why do you need to_ref() here?

I want it to evaluate to a temporary -- switched to an .eval().

SteveBronder · 2020-11-20T21:49:41Z

stan/math/rev/functor/apply_vector_unary.hpp

+/**
+ * Specialisation for use with var_value<T> types where T inherits from
+ * EigenBase. Inputs are mapped to Eigen column vectors.
+ *
+ * The returned scalar type is deduced to allow for cases where the input and
+ * return scalar types differ (e.g., functions implicitly promoting
+ * integers).
+ */


Double check these docs

SteveBronder · 2020-11-20T21:59:58Z

test/unit/math/test_ad.hpp

+template <typename ResultMatVar, typename ResultVarMat, typename MatVar,
+          typename VarMat,
+          require_std_vector_vt<is_var, ResultMatVar>* = nullptr,
+          require_std_vector_vt<is_var, ResultVarMat>* = nullptr>
+inline void test_matvar_gradient(const ad_tolerances& tols,
+                                 ResultMatVar& A_mv_f, ResultVarMat& A_vm_f,
+                                 const MatVar& A_mv, const VarMat& A_vm) {
+  for (size_t i = 0; i < A_vm_f.size(); ++i) {
+    A_vm_f[i].adj() = 1;
+    A_mv_f[i].adj() = 1;
+    stan::math::grad();
+    expect_near_rel_var("var<Matrix> vs Matrix<var> input", A_vm, A_mv, tols);
+    stan::math::set_zero_all_adjoints();
+  }
+}


This is confusing me, the requires look like it's for std::vector<var>? I think that's the only way that A_vm_f[i].adj() = 1; would work because you can't assign a constant to an entire eigen matrix like that

https://godbolt.org/z/8KnPr1

Or is this for when a var<mat> function would return a std::vector<var>?

This is for std::vector<var> outputs, not std::vector<var<mat>>

SteveBronder · 2020-11-20T22:21:28Z

test/unit/math/test_ad.hpp

+          require_std_vector_vt<is_matrix, Type>* = nullptr>
+void check_return_type(const ReturnType& ret, const Type& x) {
+  if (ret.size() > 0 && x.size() > 0)
+    check_return_type(ret[0], x[0]);


I think the logic is fine here unless I'm missing something. This is a specialization for std::vector<T> that just checks that the inner type for the input/output is correct. If that's what this does then I think it's fine

andrjohns · 2020-11-23T02:05:14Z

Quick Q - Why the addition of _impl functions here? I thought I had the apply_vector_unary implementations working, or am I remembering something else?

bbbales2 · 2020-11-23T02:22:44Z

@andrjohns it's separate implementations for var<mat> and mat<var> types. Does that make sense?

I could make them sd overloads and then specialize the apply_vector_unary when the input is a std::vector. Now that I say that outloud I kinda like it better, though it subverts some of the apply_vector_unary pattern. Also now that you mention the _impl s I guess some of the lambdas are defunct now.

andrjohns · 2020-11-23T02:33:26Z

Is the separate implementations because mat<var> isn't compatible with the calback_vari approach or because of apply_vector_unary?

Found the branch I was thinking of, I thought I had both var<mat> and mat<var> passing tests on this branch: develop...andrjohns:feature/issue-2098-vec_unary_var_mat

andrjohns · 2020-11-23T02:40:33Z

Also, sorry if this is retreading obvious stuff with var<mat>, still catching up!

bbbales2 · 2020-11-23T13:59:22Z

@andrjohns yeah this pull has those changes too (so that apply_vector_unary works with var<mat>). At least I hope they're the same only very roughly checking. I didn't realize you had a branch or I would have used that -- apologies for the duplication.

It's possible to do the mat<var> and var<mat> implementations in the same code, but two things have made me stop trying to do this:

We got gridlocked before trying to carefully benchmark mat<var> at the expense of ever getting var<mat> stuff in, so I'm just avoiding the mat<var> implementations now (it's easy to accidentally slow down the existing code by 10% and hard to get it benchmarked and fixed).
mat<var> and var<mat>, even when we do them both with reverse_pass_callback they end up looking slightly different. (.val() isn't a problem for var<mat>, but with mat<var> you want to only do it once cuz it's slow)

andrjohns · 2020-11-23T14:59:36Z

Ah that all makes sense thanks

got rid of _impl versions of sd, log_softmax, and log_sum_exp (Issue #2101)

…4.1 (tags/RELEASE_600/final)

bbbales2 · 2020-11-24T16:27:59Z

@SteveBronder I made a bunch of changes:

Added val(), adj() to all the varis. Let me know if you want this moved to a different pull or you want me to add tests. I just replaced all the .val_ and .adj_ calls in the current tests with .val() and .adj().
Added make_callback_var
I got rid of the _impl s and instead lined those up as overloads to go along with apply_vector_unary. So for each of sd, log_softmax, and log_sum_exp there's one version of the function that handles mat<var>, one version that handles var<mat>, and then an apply_vector_unary version that handles std::vector<T> (if T is a std::vector<var> apply_vector_unary changes it into an Eigen::Map). @andrjohns feel free to comment on this if you want.

SteveBronder

This all looks good to me! You have to fix up the docs but then ping me and I'll approve

SteveBronder · 2020-11-24T17:31:01Z

stan/math/rev/fun/sd.hpp

+  auto arena_diff = to_arena((x.val().array() - x.val().mean()).matrix());
+  double sum_of_squares = arena_diff.squaredNorm();
+  double sd = std::sqrt(sum_of_squares / (x.size() - 1));


[optional] You could do the little loop thing here to make this faster but it's fine as is

You have to fix up the docs

Do you mean the doxygen docs or the function reference docs (second definitely needs updated and I wouldn't doubt the first lol)?

Oh doxygen my bad

@SteveBronder I don't know what was going on with the docs. I changed all the variables to be named x to get them to work. That was a real hairtugger

…into feature/sd_varmats

…4.1 (tags/RELEASE_600/final)

bbbales2 · 2020-11-24T23:04:08Z

Woof, I had to add .val() and .adj() (and the _op) accessors to the opencl stuff. I'm firmly into I-don't-know-what-I'm-doing territory, so in the likely case this fails, I think I'll just revert the .val() and .adj() stuff.

stan-buildbot · 2020-11-25T20:18:08Z

Name	Old Result	New Result	Ratio	Performance change( 1 - new / old )
gp_pois_regr/gp_pois_regr.stan	3.63	3.59	1.01	1.24% faster
low_dim_corr_gauss/low_dim_corr_gauss.stan	0.02	0.02	1.04	3.95% faster
eight_schools/eight_schools.stan	0.12	0.11	1.04	3.47% faster
gp_regr/gp_regr.stan	0.16	0.17	0.98	-1.54% slower
irt_2pl/irt_2pl.stan	5.65	5.7	0.99	-0.97% slower
performance.compilation	86.93	85.73	1.01	1.39% faster
low_dim_gauss_mix_collapse/low_dim_gauss_mix_collapse.stan	8.43	8.54	0.99	-1.21% slower
pkpd/one_comp_mm_elim_abs.stan	28.8	29.32	0.98	-1.8% slower
sir/sir.stan	134.89	135.22	1.0	-0.25% slower
gp_regr/gen_gp_data.stan	0.04	0.04	1.0	-0.12% slower
low_dim_gauss_mix/low_dim_gauss_mix.stan	2.95	2.97	0.99	-0.68% slower
pkpd/sim_one_comp_mm_elim_abs.stan	0.38	0.37	1.01	0.53% faster
arK/arK.stan	2.48	2.51	0.99	-1.12% slower
arma/arma.stan	0.61	0.6	1.01	1.42% faster
garch/garch.stan	0.74	0.74	1.0	0.04% faster
Mean result: 1.00318155097

Jenkins Console Log
Blue Ocean
Commit hash: 422c29a

Machine information

ProductName: Mac OS X ProductVersion: 10.11.6 BuildVersion: 15G22010

CPU:
Intel(R) Xeon(R) CPU E5-1680 v2 @ 3.00GHz

G++:
Configured with: --prefix=/Applications/Xcode.app/Contents/Developer/usr --with-gxx-include-dir=/usr/include/c++/4.2.1
Apple LLVM version 7.0.2 (clang-700.1.81)
Target: x86_64-apple-darwin15.6.0
Thread model: posix

Clang:
Apple LLVM version 7.0.2 (clang-700.1.81)
Target: x86_64-apple-darwin15.6.0
Thread model: posix

Initial commit to get sd var<mat> compatible (Issue #2101)

d150183

bbbales2 marked this pull request as draft October 27, 2020 20:03

[Jenkins] auto-formatting by clang-format version 6.0.0-1ubuntu2~16.0…

bcaf552

…4.1 (tags/RELEASE_600/final)

bbbales2 commented Oct 27, 2020

View reviewed changes

bbbales2 mentioned this pull request Oct 27, 2020

Make functions with custom autodiff var<mat> friendly #2101

Open

78 tasks

bbbales2 and others added 7 commits November 17, 2020 11:22

Merge remote-tracking branch 'origin/develop' into feature/sd_varmats

a1359e7

Added testing for std::vector<var<mat>> and did more vectorizing of…

d76ad61

… `sd` (Issue #2101)

[Jenkins] auto-formatting by clang-format version 6.0.0-1ubuntu2~16.0…

6f70a44

…4.1 (tags/RELEASE_600/final)

Updated sd tests (Issue #2101)

918f1c1

Merge branch 'feature/sd_varmats' of https://github.com/stan-dev/math …

9ce0e7c

…into feature/sd_varmats

[Jenkins] auto-formatting by clang-format version 6.0.0-1ubuntu2~16.0…

f7595ab

…4.1 (tags/RELEASE_600/final)

Uncommented code (Issue #2101)

a7c9625

bbbales2 commented Nov 17, 2020

View reviewed changes

stan/math/prim/fun/sd.hpp Show resolved Hide resolved

test/unit/math/test_ad.hpp Show resolved Hide resolved

SteveBronder requested changes Nov 17, 2020

View reviewed changes

Updated log_sum_exp and log_softmax to work with std::vectors of var<…

ff99e31

…mat> types (Issue #2101)

bbbales2 commented Nov 19, 2020

View reviewed changes

yashikno and others added 2 commits November 19, 2020 20:48

Merge commit '7b9b75f8b2c897524ff522765ccad77eff99d217' into HEAD

a15f89d

[Jenkins] auto-formatting by clang-format version 6.0.0-1ubuntu2~16.0…

a8d6a5d

…4.1 (tags/RELEASE_600/final)

bbbales2 commented Nov 19, 2020

View reviewed changes

bbbales2 marked this pull request as ready for review November 19, 2020 20:58

bbbales2 changed the title ~~Initial commit to get sd var<mat> compatible (Issue #2101)~~ Get sd, log_softmax, and log_sum_exp fully var<mat> compatible (Issue #2101) Nov 19, 2020

bbbales2 added 2 commits November 20, 2020 13:26

Updated docs (Issue #2101)

13226b1

Updated two argument check_return_type to work with one argument no…

80c5f6b

…t a matrix (Issue #2101)

SteveBronder requested changes Nov 20, 2020

View reviewed changes

bbbales2 and others added 3 commits November 24, 2020 11:13

Added make_callback_var, added .val() and .adj() for vari types,

25a4b7d

got rid of _impl versions of sd, log_softmax, and log_sum_exp (Issue #2101)

Merge commit '92fce0218c9fb15fd405ef031f488cad05c5546b' into HEAD

ee3f4f7

[Jenkins] auto-formatting by clang-format version 6.0.0-1ubuntu2~16.0…

66bdb89

…4.1 (tags/RELEASE_600/final)

SteveBronder previously approved these changes Nov 24, 2020

View reviewed changes

bbbales2 added 2 commits November 24, 2020 15:27

Updated log_softmax docs (Issue #2101)

9c907c8

Merge branch 'feature/sd_varmats' of https://github.com/stan-dev/math …

ef1570d

…into feature/sd_varmats

bbbales2 dismissed SteveBronder’s stale review via ef1570d November 24, 2020 20:33

stan-buildbot and others added 3 commits November 24, 2020 20:39

[Jenkins] auto-formatting by clang-format version 6.0.0-1ubuntu2~16.0…

2ef4d87

…4.1 (tags/RELEASE_600/final)

Updated apply_vector_unary includes (Issue #2101)

dbb8935

Added accessors to opencl vari (Issue #2101)

4ec7469

Removed unnecessary test in divide (Issue #2101)

422c29a

SteveBronder approved these changes Nov 25, 2020

View reviewed changes

bbbales2 merged commit 634fc54 into develop Nov 29, 2020

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Get `sd`, `log_softmax`, and `log_sum_exp` fully `var<mat>` compatible (Issue #2101) #2169

Get `sd`, `log_softmax`, and `log_sum_exp` fully `var<mat>` compatible (Issue #2101) #2169

bbbales2 commented Oct 27, 2020 •

edited

Loading

bbbales2 left a comment

bbbales2 Oct 27, 2020

SteveBronder commented Oct 28, 2020

bbbales2 commented Nov 2, 2020

bbbales2 left a comment

SteveBronder left a comment

bbbales2 Nov 19, 2020 •

edited

Loading

bbbales2 Nov 19, 2020

bbbales2 Nov 19, 2020

SteveBronder Nov 20, 2020

SteveBronder left a comment

SteveBronder Nov 20, 2020

bbbales2 Nov 24, 2020

SteveBronder Nov 20, 2020

bbbales2 Nov 24, 2020

SteveBronder Nov 20, 2020

SteveBronder Nov 20, 2020

SteveBronder Nov 20, 2020

bbbales2 Nov 24, 2020

SteveBronder Nov 20, 2020

andrjohns commented Nov 23, 2020

bbbales2 commented Nov 23, 2020

andrjohns commented Nov 23, 2020

andrjohns commented Nov 23, 2020

bbbales2 commented Nov 23, 2020

andrjohns commented Nov 23, 2020

bbbales2 commented Nov 24, 2020

SteveBronder left a comment

SteveBronder Nov 24, 2020

bbbales2 Nov 24, 2020

bbbales2 Nov 24, 2020

bbbales2 Nov 24, 2020

bbbales2 commented Nov 24, 2020

stan-buildbot commented Nov 25, 2020

Get sd, log_softmax, and log_sum_exp fully var<mat> compatible (Issue #2101) #2169

Get sd, log_softmax, and log_sum_exp fully var<mat> compatible (Issue #2101) #2169

Conversation

bbbales2 commented Oct 27, 2020 • edited Loading

Summary

Release notes

Checklist

bbbales2 left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

SteveBronder commented Oct 28, 2020

bbbales2 commented Nov 2, 2020

bbbales2 left a comment

Choose a reason for hiding this comment

SteveBronder left a comment

Choose a reason for hiding this comment

bbbales2 Nov 19, 2020 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

SteveBronder left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

andrjohns commented Nov 23, 2020

bbbales2 commented Nov 23, 2020

andrjohns commented Nov 23, 2020

andrjohns commented Nov 23, 2020

bbbales2 commented Nov 23, 2020

andrjohns commented Nov 23, 2020

bbbales2 commented Nov 24, 2020

SteveBronder left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

bbbales2 commented Nov 24, 2020

stan-buildbot commented Nov 25, 2020

Get `sd`, `log_softmax`, and `log_sum_exp` fully `var<mat>` compatible (Issue #2101) #2169

Get `sd`, `log_softmax`, and `log_sum_exp` fully `var<mat>` compatible (Issue #2101) #2169

bbbales2 commented Oct 27, 2020 •

edited

Loading

bbbales2 Nov 19, 2020 •

edited

Loading