Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Introduce MutableTangent #626

Merged
merged 36 commits into from
Jan 25, 2024
Merged

Introduce MutableTangent #626

merged 36 commits into from
Jan 25, 2024

Conversation

oxinabox
Copy link
Member

@oxinabox oxinabox commented Aug 4, 2023

Will close #105

  • Move functionality to StructuralTangent
  • Add MutableTangent
  • Make co-PR to ChainRules.jl to add setindex
  • Make more optimal re: stable backing type

@oxinabox oxinabox marked this pull request as draft August 4, 2023 10:36
src/tangent_types/structural_tangent.jl Outdated Show resolved Hide resolved
src/tangent_types/structural_tangent.jl Outdated Show resolved Hide resolved
src/tangent_types/structural_tangent.jl Outdated Show resolved Hide resolved
src/tangent_types/structural_tangent.jl Outdated Show resolved Hide resolved
src/tangent_types/structural_tangent.jl Show resolved Hide resolved
src/tangent_types/structural_tangent.jl Outdated Show resolved Hide resolved
src/tangent_types/structural_tangent.jl Outdated Show resolved Hide resolved
src/tangent_types/structural_tangent.jl Outdated Show resolved Hide resolved
@codecov
Copy link

codecov bot commented Aug 4, 2023

Codecov Report

Attention: 7 lines in your changes are missing coverage. Please review.

Comparison is base (6664e8f) 93.91% compared to head (95e63d0) 94.15%.

❗ Current head 95e63d0 differs from pull request most recent head 73b7508. Consider uploading reports for the commit 73b7508 to get more accurate results

Files Patch % Lines
src/tangent_types/structural_tangent.jl 94.87% 6 Missing ⚠️
src/tangent_types/abstract_zero.jl 97.22% 1 Missing ⚠️
Additional details and impacted files
@@            Coverage Diff             @@
##             main     #626      +/-   ##
==========================================
+ Coverage   93.91%   94.15%   +0.24%     
==========================================
  Files          15       15              
  Lines         904      976      +72     
==========================================
+ Hits          849      919      +70     
- Misses         55       57       +2     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

src/tangent_types/structural_tangent.jl Outdated Show resolved Hide resolved
src/tangent_types/structural_tangent.jl Show resolved Hide resolved
src/tangent_types/structural_tangent.jl Outdated Show resolved Hide resolved
src/tangent_types/structural_tangent.jl Outdated Show resolved Hide resolved
src/tangent_types/structural_tangent.jl Outdated Show resolved Hide resolved
src/tangent_types/structural_tangent.jl Outdated Show resolved Hide resolved
src/tangent_types/structural_tangent.jl Outdated Show resolved Hide resolved
src/ChainRulesCore.jl Outdated Show resolved Hide resolved
src/tangent_types/structural_tangent.jl Show resolved Hide resolved
src/tangent_types/structural_tangent.jl Outdated Show resolved Hide resolved
test/tangent_types/structural_tangent.jl Outdated Show resolved Hide resolved
test/tangent_types/structural_tangent.jl Outdated Show resolved Hide resolved
test/tangent_types/structural_tangent.jl Show resolved Hide resolved
src/tangent_types/abstract_zero.jl Outdated Show resolved Hide resolved
src/tangent_types/structural_tangent.jl Show resolved Hide resolved
src/tangent_types/structural_tangent.jl Outdated Show resolved Hide resolved
src/tangent_types/structural_tangent.jl Show resolved Hide resolved
src/tangent_types/structural_tangent.jl Outdated Show resolved Hide resolved
test/tangent_types/abstract_zero.jl Show resolved Hide resolved
test/tangent_types/structural_tangent.jl Outdated Show resolved Hide resolved
test/tangent_types/structural_tangent.jl Outdated Show resolved Hide resolved
test/tangent_types/structural_tangent.jl Outdated Show resolved Hide resolved
test/tangent_types/structural_tangent.jl Outdated Show resolved Hide resolved
Copy link
Member

@oscardssmith oscardssmith left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is a pretty big change, but It looks good.

Copy link
Member

@willtebbutt willtebbutt left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is cool. Glad to see it added. I've a few things to say about zero_tangent because I'm implemented something similar recently.

src/tangent_types/abstract_zero.jl Outdated Show resolved Hide resolved
src/tangent_types/abstract_zero.jl Outdated Show resolved Hide resolved
src/tangent_types/abstract_zero.jl Outdated Show resolved Hide resolved
src/tangent_types/abstract_zero.jl Outdated Show resolved Hide resolved
test/tangent_types/abstract_zero.jl Show resolved Hide resolved
test/tangent_types/structural_tangent.jl Show resolved Hide resolved
src/tangent_types/abstract_zero.jl Outdated Show resolved Hide resolved
src/tangent_types/abstract_zero.jl Outdated Show resolved Hide resolved
src/tangent_types/abstract_zero.jl Outdated Show resolved Hide resolved
src/tangent_types/abstract_zero.jl Outdated Show resolved Hide resolved
test/tangent_types/abstract_zero.jl Outdated Show resolved Hide resolved
test/tangent_types/abstract_zero.jl Show resolved Hide resolved
test/tangent_types/abstract_zero.jl Show resolved Hide resolved
test/tangent_types/abstract_zero.jl Outdated Show resolved Hide resolved
test/tangent_types/abstract_zero.jl Outdated Show resolved Hide resolved
return Array{guess_zero_tangent_type(T),N}
end
guess_zero_tangent_type(::Any) = Any # if we had a general way to handle determining tangent type # https://github.com/JuliaDiff/ChainRulesCore.jl/issues/634
# TODO: we might be able to do better than this. even without.
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[JuliaFormatter] reported by reviewdog 🐶

Suggested change
# TODO: we might be able to do better than this. even without.
# TODO: we might be able to do better than this. even without.

@willtebbutt
Copy link
Member

It might make sense to test that this works for self-referential types c.f. something like this example from the docs. Maybe add a differentiable field of some kind? I'm not entirely sure how this will work -- I'm still trying to figure out how to avoid an infinite-loop in some of my work.

@ToucheSir
Copy link
Contributor

I don't see any mention of self-referential types in the Swift docs on automatic tangent type synthesis, but I imagine they must've run into this so perhaps the implementation could have some clues?

src/tangent_types/abstract_zero.jl Outdated Show resolved Hide resolved
src/tangent_types/abstract_zero.jl Outdated Show resolved Hide resolved
src/tangent_types/abstract_zero.jl Outdated Show resolved Hide resolved
src/tangent_types/abstract_zero.jl Outdated Show resolved Hide resolved
src/tangent_types/structural_tangent.jl Outdated Show resolved Hide resolved
test/tangent_types/abstract_zero.jl Outdated Show resolved Hide resolved
test/tangent_types/abstract_zero.jl Outdated Show resolved Hide resolved
test/tangent_types/abstract_zero.jl Outdated Show resolved Hide resolved
test/tangent_types/abstract_zero.jl Outdated Show resolved Hide resolved
test/tangent_types/abstract_zero.jl Outdated Show resolved Hide resolved
@oxinabox
Copy link
Member Author

Ok this has now been massively overhauled.
Mostly in the zero_tangent code.
Special attention has been paid to achieve type stability where possibly,
and to have correct behavour where the primal has nonconcretely typed fields or where it has
With the intent that we can also use zero_tangent rather than ZeroTangent when we need type stable for literals or other other times we get it from frules if that is the AD's preference.
Doing this kinda sucked but I think i worked all the edge cases out.

However, in the process some things may have been broken.
So the next step will be to fix all those and get CI passing.

I expect MutableTangent will be missing some of the nicer overloads for linear operators etc.,
but those can come later as we find they are needed.

@oxinabox
Copy link
Member Author

zero_tangent should be more or less idempotent.
Anything it returns should also be a valid zero for itself. More or less.
We should have tests for this.

This property, in general, holds for all valid tangent types.
We should document this.

@oscardssmith
Copy link
Member

This looks reasonable to me.

@oxinabox
Copy link
Member Author

The error in the projection tests is unrelated to this PR and is fixed in
#645 (by deleting those tests)

oxinabox and others added 23 commits January 25, 2024 17:44
Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
Copy link
Member

@willtebbutt willtebbutt left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is cool.

@oxinabox oxinabox merged commit 96e0f54 into main Jan 25, 2024
17 of 26 checks passed
@oxinabox oxinabox deleted the ox/mutabletangent branch January 25, 2024 11:05
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

MutableTangent
4 participants