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

RuleConfigs (include for calling back into AD) #363

Merged
merged 17 commits into from
Jun 11, 2021
Merged

RuleConfigs (include for calling back into AD) #363

merged 17 commits into from
Jun 11, 2021

Conversation

oxinabox
Copy link
Member

@oxinabox oxinabox commented Jun 1, 2021

Closes #68 and I think another one about adding config obects that I can't find

May prove useful for #350 if we make rules use a special feature to say they don't supprt inplact accumulation and then we use this to disable the copying.

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.

Excited to see where this goes.

src/rules.jl Outdated Show resolved Hide resolved
src/rules.jl Outdated Show resolved Hide resolved
src/rules.jl Outdated Show resolved Hide resolved
docs/src/config.md Show resolved Hide resolved
docs/src/config.md Outdated Show resolved Hide resolved
Copy link
Member

@mzgubic mzgubic left a comment

Choose a reason for hiding this comment

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

As far as I understand, this is not a breaking change, right? The rules can still be written in the same way (without the config). The AD packages can still call the rules in the same way (without providing the config).

docs/src/config.md Outdated Show resolved Hide resolved
docs/src/config.md Outdated Show resolved Hide resolved
docs/src/config.md Outdated Show resolved Hide resolved
docs/src/config.md Outdated Show resolved Hide resolved
docs/src/config.md Outdated Show resolved Hide resolved
docs/src/config.md Outdated Show resolved Hide resolved
src/rules.jl Outdated Show resolved Hide resolved
src/rules.jl Outdated Show resolved Hide resolved
src/rules.jl Outdated Show resolved Hide resolved
src/rules.jl Outdated Show resolved Hide resolved
@oxinabox
Copy link
Member Author

oxinabox commented Jun 2, 2021

As far as I understand, this is not a breaking change, right? The rules can still be written in the same way (without the config). The AD packages can still call the rules in the same way (without providing the config).

Correct. At least in its current form. but still WIP

@oxinabox
Copy link
Member Author

oxinabox commented Jun 7, 2021

TODO: check that this works with zygote via makign a PR to use a rule via branching from FluxML/Zygote.jl#987

@codecov-commenter
Copy link

codecov-commenter commented Jun 7, 2021

Codecov Report

Merging #363 (e44cee2) into master (2d64bce) will decrease coverage by 1.84%.
The diff coverage is 100.00%.

Impacted file tree graph

@@            Coverage Diff             @@
##           master     #363      +/-   ##
==========================================
- Coverage   89.26%   87.42%   -1.85%     
==========================================
  Files          12       14       +2     
  Lines         475      485      +10     
==========================================
  Hits          424      424              
- Misses         51       61      +10     
Impacted Files Coverage Δ
src/ChainRulesCore.jl 100.00% <ø> (ø)
src/config.jl 100.00% <100.00%> (ø)
src/rules.jl 100.00% <100.00%> (ø)
src/differentials/abstract_differential.jl 50.00% <0.00%> (-16.67%) ⬇️
src/differentials/abstract_zero.jl 92.85% <0.00%> (-1.27%) ⬇️
src/differentials/notimplemented.jl 70.83% <0.00%> (-1.17%) ⬇️
src/differentials/composite.jl 82.90% <0.00%> (-0.15%) ⬇️
src/deprecated.jl 0.00% <0.00%> (ø)
src/accumulation.jl 97.22% <0.00%> (+0.07%) ⬆️
... and 1 more

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update 2d64bce...e44cee2. Read the comment docs.

@oxinabox
Copy link
Member Author

oxinabox commented Jun 8, 2021

  • RuleConfig's need to broadcast like scalars

docs/src/config.md Outdated Show resolved Hide resolved
docs/src/config.md Outdated Show resolved Hide resolved
docs/src/config.md Outdated Show resolved Hide resolved
docs/src/config.md Show resolved Hide resolved
docs/src/config.md Show resolved Hide resolved
src/rules.jl Outdated Show resolved Hide resolved
test/config.jl Outdated Show resolved Hide resolved
test/config.jl Outdated Show resolved Hide resolved
test/config.jl Outdated Show resolved Hide resolved
test/config.jl Outdated Show resolved Hide resolved
Copy link
Member

@mzgubic mzgubic left a comment

Choose a reason for hiding this comment

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

My understanding/opinion is that

  1. traits are defined in ChainRulesCore
  2. each AD package defines exactly one config (and no traits)
  3. rule authors do not define configs nor traits

if AD defines traits, then rule authors have to depend on the AD system to write rules for it. Or the AD system has to depend on the package to define a rule for it. Neither options are great.

Rule authors have no business defining configs. They possibly could define traits, but that would need AD would have to depend on that package. Which seems bad.

On the other hand, defining an OnlyOnZygote trait inside ChainRulesCore seems bad and an implicit reverse dependency. Possibly this means we should not have OnlyOnZygote-like traits at all. The traits should only be about properties of AD systems. If that means having a BadWithLoops trait, then we'll take that over OnlyOnZygote trait.

test/config.jl Show resolved Hide resolved
test/config.jl Outdated Show resolved Hide resolved
test/config.jl Outdated Show resolved Hide resolved
test/config.jl Outdated Show resolved Hide resolved
@oxinabox
Copy link
Member Author

My understanding/opinion is that

1. traits are defined in ChainRulesCore

2. each AD package defines exactly one config (and no traits)

3. rule authors do not define configs nor traits

if AD defines traits, then rule authors have to depend on the AD system to write rules for it. Or the AD system has to depend on the package to define a rule for it. Neither options are great.

Rule authors have no business defining configs. They possibly could define traits, but that would need AD would have to depend on that package. Which seems bad.

Basically, yes.
Potentially, rule authors might define traits, as you say.
But then (as you say) the AD authour would have to depend on the package definine the rules.
Or some how allow the end-user to add traits to its config, which I can imagine being a thing.

(One use case I can speculate for Rule authors defining traits is resolving type-piracy: if DistributionsAD.jl, and DistributionsDiff.jl both provide rules for Distributions.jl, then they could do so by defining a trait that the rule config needs to have)

On the other hand, defining an OnlyOnZygote trait inside ChainRulesCore seems bad and an implicit reverse dependency. Possibly this means we should not have OnlyOnZygote-like traits at all.

Yes, we don't have a OnlyOnZygote trait.
Either we have a HatesLoopsTrait defined in ChainRulesCore that rule authors can use.
Or for Zygote only things (like the example that modifies add!!) rules can be defines in Zygote itself that dispatch on the ZygoteRuleConfig (not on a trait it is)

@oxinabox oxinabox changed the title WIP: Initial sketch for RuleConfigs (include for calling back into AD) RuleConfigs (include for calling back into AD) Jun 10, 2021
@oxinabox
Copy link
Member Author

This breaks ChainRulesOverloadGeneration.jl
I kinda don't care, cos no one is using it, but probably I should work why (so we don't get a bunch of alerts) and fix it in ChaunRulesOverloadGeneration.jl
I don't think it will be possible to not break it here, because that package cares about the internals of ChainRulesCore.
(kinda why it was in the package to begin with)

docs/src/config.md Outdated Show resolved Hide resolved
docs/src/config.md Outdated Show resolved Hide resolved
src/rules.jl Outdated Show resolved Hide resolved
@oschulz
Copy link
Contributor

oschulz commented Jun 11, 2021

Awesome, I had been looking forward to this (have some cases where I currently have to call back into Zygote explicitly).

Thanks!

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.

rrule_via_ad // frule_via_ad Calling back into AD from ChainRules
5 participants