Remove Composite Bundle #216
Conversation
|
|
||
| @Base.constprop :aggressive function (::∂☆{N})(::ATB{N, typeof(getfield)}, x::CompositeBundle{N, B}, s::AbstractTangentBundle{N, Symbol}) where {N, B} | ||
| x.tup[Base.fieldindex(B, primal(s))] | ||
| end |
There was a problem hiding this comment.
this shouldn't be needed as we already have cases for TaylorBundle above.
Likewise several others of this kind below will not be needed.
|
|
||
| let var"'" = Diffractor.PrimeDerivativeFwd | ||
| @test foo'(100.0) == 2.0 | ||
| @test foo''(100.0) == 0.0 |
There was a problem hiding this comment.
this second derivative is actually broken on main
There was a problem hiding this comment.
do we know what change broke it?
|
|
||
| let var"'" = Diffractor.PrimeDerivativeFwd | ||
| @test foo'(100.0) == 2.0 | ||
| @test foo''(100.0) == 0.0 |
There was a problem hiding this comment.
this second derivative is actually broken on main
Codecov ReportAttention:
Additional details and impacted files@@ Coverage Diff @@
## main #216 +/- ##
==========================================
+ Coverage 54.90% 55.05% +0.14%
==========================================
Files 28 28
Lines 2801 2810 +9
==========================================
+ Hits 1538 1547 +9
Misses 1263 1263
☔ View full report in Codecov by Sentry. |
| _TangentBundle(Val{N}(), primal, TaylorTangent(coeffs)) | ||
| end | ||
|
|
||
| function check_taylor_invariants(coeffs, primal, N) | ||
| @assert length(coeffs) == N | ||
| if isa(primal, TangentBundle) | ||
| @assert isa(coeffs[1], TangentBundle) |
There was a problem hiding this comment.
this check is just wrong AFAICT.
because if the primal is a TangentBundle, then the first derivative is a Tangent{TangentBundle}
The first deriviative has no primal component.
But this was never hit before because we were not making TaylorBundles when taking derivatives of TaylorBundles, (at least not in anything that happens in our test/forward.jl I checked) we were (I assume) making CompositeBundles
| a_partials = ntuple(i->partial(r, ii)[1], N) | ||
| b_partials = ntuple(i->partial(r, ii)[2], N) | ||
| the_partials = (a_partials..., primal_b, b_partials...) | ||
| return TangentBundle{N+1}(the_primal, the_partials) | ||
| end | ||
| end |
There was a problem hiding this comment.
For reference as a part way step to determining this , I first refactored the CompositeBundle version into
function shuffle_up(r::CompositeBundle{N}) where {N}
a, b = r.tup
if isswifty(a) && isswifty(b) && taylor_compatible(a, b)
the_partials = ntuple(N+1) do i
if ii <= N
a[TaylorTangentIndex(i)] # == b[TaylorTangentIndex(i-1)] (except first which is b.primal)
else # ii = N+1
b[TaylorTangentIndex(i-1)]
end
end
return TaylorBundle{N+1}(primal(a), the_partials)
else
the_primal = r.tup[1].primal
a_partials = r.tup[1].tangent.partials
b_partials = ntuple(i->partial(b,i), 1<<(N+1)-1)
the_partials = (a_partials..., primal_b, b_partials...)
return TangentBundle{N+1}(the_primal, the_partials)
end
enduse fieldcount taylor bundle not taylor tangent
oxinabox
force-pushed
the
ox/nocompo
branch
from
93d66c2 to
bfd99c2
Compare
| "for a TaylorTangent{N, <:Tuple} this breaks it up unto 1 TaylorTangent{N} for each element of the primal tuple" | ||
| function destructure(r::TaylorBundle{N, B}) where {N, B<:Tuple} | ||
| return ntuple(fieldcount(B)) do field_ii | ||
| the_primal = primal(r)[field_ii] | ||
| the_partials = ntuple(N) do order_ii | ||
| partial(r, order_ii)[field_ii] | ||
| end | ||
| return TaylorBundle{N}(the_primal, the_partials) | ||
| end | ||
| end |
There was a problem hiding this comment.
Do we need this, or can we just use unbundle?
src/stage1/forward.jl
Outdated
| a_partials = ntuple(i->partial(r, ii)[1], N) | ||
| b_partials = ntuple(i->partial(r, ii)[2], N) | ||
| the_partials = (a_partials..., primal_b, b_partials...) | ||
| return TangentBundle{N+1}(the_primal, the_partials) |
There was a problem hiding this comment.
my problem with this is I don;t think that b_partials length 1<<(N+1)-1
because it has length N
and those are never equal
|
Nightly is failing just cos of reverse mode, unreleated to this PR |
|
I don't exactly remember what the original motivation for this was, but if everything goes through fine, I have no particular attachment to this type. |
Co-authored-by: Elliot Saba <[email protected]>
This is a prerequisite for #189 as how CompositeBundle works, causes a problem when it comes to mutation.
Right now
CompositeBundlestores everything as a tuple of other Bundles, one per field of the primal.Rather than storing things as a a a primal and a tangent (or really 1 or more tangents).
This is a problem because when we need to get stuff back seperated into primal and tangent (which we need to do to apply
frulesor to actually get at the derivative result) we need to assemble the Tangent (note the capital T, this is the type that represents an immutable structural tangent) on the fly.Which we do here:
Diffractor.jl/src/stage1/forward.jl
Lines 8 to 14 in ee8af85
That is all well and good if everything is immuatable, since there is no need to preserve mutation or worry about aliasing if things are immuable.
If on the other hand, things are mutable, you do need to worry about this things. In particular the
MutableTangenttype must be aliased in the same places its primal is (infact in the PR JuliaDiff/ChainRulesCore.jl#626 that created that, I always made a note in the docstring), you very much can not break it a part and recreate it on demand.Ergo, I think CompositeBundle needs to be removed and replaced with something that works like the other tangent types, a pair of primal and tangent values.
Namely
TaylorTangentin most cases, though in theory it might end up with anExplictTangentHowever, actually removing
CompositeBundleseems a net significant simplification of the code base.This PR is net negative in lines of code in
src/and will get even more negative once I work out which of the remaining const prop aggressive special cases can be deleted.It also fixes some existing issues and I think will make others easier to fix in the future.