Fix with_logabsdet_jacobian for mapped/broadcasted functions - #40
Fix with_logabsdet_jacobian for mapped/broadcasted functions#40oschulz wants to merge 5 commits into
Conversation
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## master #40 +/- ##
==========================================
+ Coverage 97.64% 97.77% +0.13%
==========================================
Files 5 5
Lines 85 90 +5
==========================================
+ Hits 83 88 +5
Misses 2 2 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
21489fe to
15abf0a
Compare
|
Any objections @devmotion ? |
15abf0a to
7732448
Compare
|
Good like this now @devmotion ? |
| function _with_ladj_on_mapped(y_with_ladj::Tuple{NoLogAbsDetJacobian, Vararg{Union{NoLogAbsDetJacobian, Tuple{Any,Real}}}}) | ||
| return first(y_with_ladj) | ||
| end | ||
|
|
||
| function _with_ladj_on_mapped(y_with_ladj::Tuple{Tuple{Any,Real}, Vararg{Union{NoLogAbsDetJacobian, Tuple{Any,Real}}}}) | ||
| a, bs = first(y_with_ladj), Base.tail(y_with_ladj) | ||
| return foldl(_combine_y_ladj, bs, init = ((a[1],), a[2])) | ||
| end |
There was a problem hiding this comment.
I wonder - can we limit the summation already on the dispatch level to the case without NoLogAbsDetJacobian, by restricting the acutal summation path to Tuple{Tuple{Any,Real}, Vararg{Tuple{Any,Real}}}? All other cases should result in NoLogAbsDetJacobian anyway? And presumably, in contrast to the current foldl result, not in any NoLogAbsDetJacobian but the one of the mapped_f and X, to be consistent with other NoLogAbsDetJacobian results?
There was a problem hiding this comment.
Tuple{Any,Real} will cause problems with Reactant, which I definitely want to use with ChangesOfVariables soon. And while it's not sanctioned I've been using with_logabsdet_jacobian in machine-learning cases that necessarily use a flat array for batching instead of nested arrays, and so the ladj value then has to be a vector, not a number. I've been thinking about making it "official", but at least for now, we only required ladj to be additive.
There was a problem hiding this comment.
@devmotion I changed the dispatch now on the lines you suggested, based on "positive" instead of "negative" patterns.
3a95955 to
cb671fb
Compare
2 participants
Occurances of
NoLogAbsDetJacobianin intermediate result were not handled properly.