Insert `_project` into `getproperty`'s gradient, and then improve `z2d` etc. to restore stability #1104

mcabbott · 2021-10-16T05:15:33Z

The goal here is to avoid things like this:

julia> pullback(x -> abs2(x.im * ((1-im) * x).re), 4+5im)[2](1.0)
# Zygote v0.6.0:
((re = 450.0 + 450.0im, im = 810.0),)
# Zygote v0.6.28:
ERROR: MethodError: no method matching Complex(::ComplexF64, ::ZeroTangent)
# this PR:
(450.0 + 1260.0im,)

While projection was intended as a mechanism to encode & enforce mathematical properties like real vs complex, I think it's also a good way to encode (when desired) a preference for the "natural gradient" of a complex number, instead of a NamedTuple, because this can participate in other stages of the calculation. One example of this found in the wild was fixed in JuliaDiff/ChainRules.jl#509, a made-up example in the same spirit is:

julia> gradient(x -> sqrt(sum(parent(x .^ 2))), Diagonal([1,2,3]))[1]
# latest tagged:
ERROR: ArgumentError: broadcasting over dictionaries and `NamedTuple`s is reserved
# this PR + CRC#446:
3×3 Diagonal{Float64, Vector{Float64}}:
 0.267261   ⋅         ⋅ 
  ⋅        0.534522   ⋅ 
  ⋅         ⋅        0.801784

I'm a little concerned this may affect what #909 just fixed. Maybe CI will tell us, but cc @simeonschaub ?

simeonschaub · 2021-10-16T05:25:24Z

Seems reasonable, since _project should be type stable I don't expect this to cause too many issues. We probably want to do this in the mutable case as well though for consistency?

mcabbott · 2021-10-16T05:28:48Z

Ok, good. Yes I need to stare a bit more at the mutable case, and think up some tests.

willtebbutt · 2021-10-16T08:59:26Z

Should we need this here? There's nothing in the default getproperty implementation that would warrant it -- if something needs projecting here, presumably it's because a rule somewhere else was incorrect, no?

mcabbott · 2021-10-16T12:48:16Z

because a rule somewhere else was incorrect, no?

No. These examples aren't doing "mathematical" projection at all. They just change the representation, from a NamedTuple to something more like the original.

As I tried to say above, this is a second meaning to overload onto _project, and there is some chance that it ought to be called something else. But I think that something else would share many features, it wants to be similarly user-overloadable. Combining these seems fairly natural and tidy to me.

default getproperty implementation

There's nothing wrong with the default, except that it returns a NamedTuple. The examples above have steps before getproperty, which the gradient wants to flow backwards into, and these steps don't work with a NamedTuple -- you get the errors shown or the nonsensical answers shown.

willtebbutt · 2021-10-18T10:09:39Z

Ahh I see, thanks for the explanation. The only basis on which I think I take issue with doing this is on aspects of how we're doing AD that I know we already disagree about (when to use naturals vs structurals etc), so I won't get in the way of this.

mcabbott

90% of the PR is upgrades & bug fixes to the map to & from ChainRules types. They aren't perfect yet, but better.

_project now uses zygote2differential and hence does not produce Tangent{Any}, so we can remove some CRC piracy which handled that.

src/compiler/chainrules.jl

src/lib/lib.jl

test/compiler.jl

test/features.jl

mcabbott · 2021-10-24T13:46:05Z

.github/workflows/Downstream.yml

+          - {user: TuringLang, repo: DynamicPPL.jl, group: All}
+          - {user: TuringLang, repo: DistributionsAD.jl, group: Zygote}


Downstream test of DistributionsAD is trivial right now. When TuringLang/DistributionsAD.jl#203 is merged, it should then test just the Zygote test group.

It looks like that should take < 30 mins. Running all tests takes > 6 hours, times out on github actions.

src/compiler/chainrules.jl

oxinabox · 2021-10-26T14:58:17Z

src/compiler/chainrules.jl

-  primals = NamedTuple{fnames}(getfield(primal, fn) for fn in fnames)
-  tp::NamedTuple = map(z2d, complete_t, primals)
-  return canonicalize(Tangent{primal_type, typeof(tp)}(tp))
+z2d(dx::NamedTuple, primal::AbstractDict) = dx  # uses a NamedTuple but not for fields!


Suggested change

z2d(dx::NamedTuple, primal::AbstractDict) = dx # uses a NamedTuple but not for fields!

# Dict handling in Zygote is a mess for now leave it alone.

# TODO should this become a `Tangent{Dict,Dict}` ?

# right now it uses a NamedTuple but not for fields of AbstractDict struict

z2d(dx::NamedTuple, primal::AbstractDict) = dx

test/features.jl

oxinabox · 2021-10-26T15:10:13Z

Does this close #660 ?

test/runtests.jl

mcabbott · 2021-10-26T15:23:20Z

Does this close #660 ?

That appears to be fixed, but not by this PR. Can add it as a test to another list of old PRs I was making...

src/compiler/chainrules.jl

oxinabox · 2021-10-29T13:31:14Z

I am going to leave this to @DhairyaLGandhi to give final approval

devmotion · 2021-10-30T19:54:25Z

FYI this PR fixes the remaining Zygote test errors in TuringLang/DistributionsAD.jl#203 (TuringLang/DistributionsAD.jl#203 (comment)).

DhairyaLGandhi · 2021-10-31T21:26:27Z

Cool, let me take a pass at this in my morning. @devmotion does this address the accum(::NamedTuple, ::Tangent) case for you?

devmotion · 2021-10-31T21:30:29Z

@devmotion does this address the accum(::NamedTuple, ::Tangent) case for you?

I don't know which error you refer to, it's none of the errors in TuringLang/DistributionsAD.jl#203 (comment). Both Zygote problems in this comment (TuringLang/DistributionsAD.jl#203 (comment)) are fixed with this PR though.

src/compiler/chainrules.jl

DhairyaLGandhi · 2021-11-01T06:42:07Z

src/compiler/chainrules.jl

-  return canonicalize(Tangent{primal_type, typeof(tp)}(tp))
+# Dict handling in Zygote is a mess... should this become a  `Tangent{Dict,Dict}` ?
+# Right now it uses a NamedTuple but not for fields of the AbstractDict struct
+z2d(dx::NamedTuple, primal::AbstractDict) = dx


Why should it be the fields of the AbstractDict struct?

No statement is being made here about should or shouldn't.
Per the comment Dict handling is currently a mess, and is inconsistent.
But most of the time it is a named tuple that is not the fields, but rather is the values.
And that is fine.

DhairyaLGandhi · 2021-11-01T06:49:53Z

src/compiler/chainrules.jl

+  if backing isa Tuple{Vararg{AbstractZero}}
+    return NoTangent()  # collapse all-zero case


There seems to be a lot of "collapsing zero" going on. The accumulation step already takes care of the collapsing for the most part. Why can't we let the zeros cases get handled in the accumulation step? Is this related to zygote sometimes producing Array{Nothing}?

Without it the DistributionsAD tests fail. The main reason is that ChainRules expects non-differentiable objects to have a NoTangent derivative but without this fix we end up with stuff like Tangent{T}(; a = NoTangent(), b = [NoTangent(), NoTangent()], c = Tangent{S}(; d = NoTangent())).

A stacktrace would be helpful.

There are more explanations in TuringLang/DistributionsAD.jl#203 (comment). Unfortunately, it seems there are no GH actions logs of the original issue since I fixed it (partially) before adding and pushing a new commit.

DhairyaLGandhi · 2021-11-01T06:53:23Z

Given the extensive breakage in the ecosystem, we would definitely need a hot patch so we can get PPL and SciML up and running again. We should look for a more robust solution than we have currently to handle the zygote <-> CR interaction.

devmotion · 2021-11-01T08:38:47Z

IMO the integration tests added in this PR are a big improvement and hopefully ensure that Zygote support in Distributions and Turing is not broken accidentally again.

I don't think it's the main problem right now but I assume the CR-Zygote integration could be simplified if Zygote would use NoTangent and ZeroTangent instead of nothing.

DhairyaLGandhi · 2021-11-02T15:20:16Z

src/compiler/chainrules.jl

+# Could `reinterpret` instead here...
+# One easy case, but can this go wrong?
+# @inline wrap_chainrules_input(xs::Base.ReinterpretArray{<:NamedTuple, <:Tangent}) = parent(@show xs)


This seems like debugging code

The show you mean? It's in a comment, about further optimisations, for which again 1112 is the issue.

We can get rid of the show either way. In general, it's best to leave reinterpreted arrays be, they can have unexpected behavior if we try to manage them manually. Even in a comment, I still doubt the show was intentionally expected to be part of committed code?

DhairyaLGandhi · 2021-11-02T15:25:14Z

src/compiler/chainrules.jl

-# Could `reinterpret` instead of broadcasting here -- TODO
-@inline wrap_chainrules_output(xs::AbstractArray{<:ChainRules.Tangent}) = wrap_chainrules_output.(xs)
+wrap_chainrules_output(dxs::AbstractArray{<:Number}) = dxs
+wrap_chainrules_output(dxs::AbstractArray{<:AbstractArray{<:Number}}) = dxs


Is there an mwe for this being needed? Is it for RecursiveArrayTools or does it come up elsewhere?

For all of these, the minimal correct thing is wrap_chainrules_output(dxs::AbstractArray) = map(wrap_chainrules_output, dxs), etc.

But for some common types, we know that the representation used by Zygote and by ChainRules agrees, so we can save a copy by doing nothing. That's the case for AbstractArray{<:Number} and AbstractArray{<:AbstractArray{<:Number}}. These aren't needed, but are an optimisation for common cases.

1112 is the issue about less trivial optimisations to save more copies.

DhairyaLGandhi

bors r+

ToucheSir · 2021-11-07T21:28:25Z

Cryptic error from Bors: https://app.bors.tech/repositories/20723/log

mcabbott · 2021-11-07T21:40:17Z

Indeed.

This did pass buildkite tests, BTW, after being rebased on master (a8306bc), and master hasn't changed since.

FluxML/Zygote.jl#1104 was merged and hence (hopefully) Zygote tests are not broken anymore.

mcabbott force-pushed the project_getprop branch 2 times, most recently from 28c05a1 to 34f6af6 Compare October 17, 2021 18:36

oschulz mentioned this pull request Oct 17, 2021

Utilize ChainRulesCore thunks #966

Open

This was referenced Oct 18, 2021

Make ProjectTo convert Tangent back to Diagonal, etc, when safe JuliaDiff/ChainRulesCore.jl#446

Closed

[Merged by Bors] - Set up documentation TuringLang/DynamicPPL.jl#334

Closed

Neural adapter test is broken SciML/NeuralPDE.jl#412

Closed

mcabbott commented Oct 21, 2021

View reviewed changes

mcabbott marked this pull request as ready for review October 21, 2021 00:17

This was referenced Oct 22, 2021

Remove adjoint for fill and fix tests TuringLang/DistributionsAD.jl#203

Merged

ChainRulesCore.NoTangent error when updating Zygote 0.6.21 -> 0.6.22 mcabbott/SliceMap.jl#12

Closed

mcabbott commented Oct 24, 2021

View reviewed changes

mcabbott requested a review from oxinabox October 24, 2021 13:52

oxinabox reviewed Oct 26, 2021

View reviewed changes

src/compiler/chainrules.jl Show resolved Hide resolved

oxinabox reviewed Oct 26, 2021

View reviewed changes

src/compiler/chainrules.jl Outdated Show resolved Hide resolved

oxinabox reviewed Oct 26, 2021

View reviewed changes

test/features.jl Show resolved Hide resolved

oxinabox reviewed Oct 26, 2021

View reviewed changes

test/runtests.jl Show resolved Hide resolved

mcabbott added 7 commits October 26, 2021 11:15

insert _project into getproperty

1f5602c

use zygote2differential in _project

69ba274

improve type-stability of zygote2differential

07aca19

fix 1 test and break 2

0c7eedc

2 not broken in fact

b8b6782

skip inference test

6634c10

skip more inference tests

7730d6d

mcabbott added 2 commits October 26, 2021 11:15

DistributionsAD too

2e9c1a0

collapse zeros in z2d

19d24a9

mcabbott added 2 commits October 26, 2021 11:25

comments

775fb13

indents

a8306bc

mcabbott force-pushed the project_getprop branch from 45f5c27 to a8306bc Compare October 26, 2021 15:25

mcabbott changed the title ~~Insert _project into getproperty's gradient~~ Insert _project into getproperty's gradient, and then improve z2d etc. to restore stability Oct 26, 2021

This was referenced Oct 27, 2021

Mutable struct bugs #1111

Open

NNPDENS and NNPDEHan tests are still failing SciML/NeuralPDE.jl#420

Closed

DhairyaLGandhi reviewed Oct 27, 2021

View reviewed changes

src/compiler/chainrules.jl Show resolved Hide resolved

src/compiler/chainrules.jl Show resolved Hide resolved

wupeifan mentioned this pull request Oct 28, 2021

Issues after derivatives hooks are added HighDimensionalEconLab/DifferentiableStateSpaceModels.jl#64

Closed

DhairyaLGandhi reviewed Nov 1, 2021

View reviewed changes

DhairyaLGandhi reviewed Nov 2, 2021

View reviewed changes

change one comment

5781c30

DhairyaLGandhi approved these changes Nov 7, 2021

View reviewed changes

mcabbott merged commit 4ed3a86 into FluxML:master Nov 7, 2021

mcabbott deleted the project_getprop branch November 7, 2021 21:15

devmotion mentioned this pull request Nov 7, 2021

[Merged by Bors] - Zygote tests are (?) not broken anymore TuringLang/DynamicPPL.jl#339

Closed

bors bot pushed a commit to TuringLang/DynamicPPL.jl that referenced this pull request Nov 8, 2021

Zygote tests are (?) not broken anymore (#339)

ad545be

FluxML/Zygote.jl#1104 was merged and hence (hopefully) Zygote tests are not broken anymore.

devmotion mentioned this pull request Nov 8, 2021

remove @adjoint function cholesky #1114

Merged

ToucheSir mentioned this pull request Mar 28, 2022

ChainRulesCore.Tangent type leaking erroing on Zygote.Jnew #1107

Open

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Insert `_project` into `getproperty`'s gradient, and then improve `z2d` etc. to restore stability #1104

Insert `_project` into `getproperty`'s gradient, and then improve `z2d` etc. to restore stability #1104

mcabbott commented Oct 16, 2021

simeonschaub commented Oct 16, 2021

mcabbott commented Oct 16, 2021

willtebbutt commented Oct 16, 2021

mcabbott commented Oct 16, 2021 •

edited

Loading

willtebbutt commented Oct 18, 2021

mcabbott left a comment •

edited

Loading

mcabbott Oct 24, 2021

oxinabox Oct 26, 2021

oxinabox commented Oct 26, 2021

mcabbott commented Oct 26, 2021 •

edited

Loading

oxinabox commented Oct 29, 2021

devmotion commented Oct 30, 2021

DhairyaLGandhi commented Oct 31, 2021

devmotion commented Oct 31, 2021

DhairyaLGandhi Nov 1, 2021

oxinabox Nov 1, 2021

DhairyaLGandhi Nov 1, 2021

devmotion Nov 1, 2021

DhairyaLGandhi Nov 1, 2021

devmotion Nov 1, 2021

DhairyaLGandhi commented Nov 1, 2021 •

edited

Loading

devmotion commented Nov 1, 2021

DhairyaLGandhi Nov 2, 2021

mcabbott Nov 3, 2021 •

edited

Loading

DhairyaLGandhi Nov 4, 2021

DhairyaLGandhi Nov 2, 2021

mcabbott Nov 3, 2021 •

edited

Loading

DhairyaLGandhi left a comment

ToucheSir commented Nov 7, 2021

mcabbott commented Nov 7, 2021

		- {user: TuringLang, repo: DynamicPPL.jl, group: All}
		- {user: TuringLang, repo: DistributionsAD.jl, group: Zygote}

-z2d(dx::NamedTuple, primal::AbstractDict) = dx  # uses a NamedTuple but not for fields!
+# Dict handling in Zygote is a mess for now leave it alone.
+# TODO should this become a  `Tangent{Dict,Dict}` ?
+# right now it uses a NamedTuple but not for fields of AbstractDict struict
+z2d(dx::NamedTuple, primal::AbstractDict) = dx

		if backing isa Tuple{Vararg{AbstractZero}}
		return NoTangent() # collapse all-zero case

Insert _project into getproperty's gradient, and then improve z2d etc. to restore stability #1104

Insert _project into getproperty's gradient, and then improve z2d etc. to restore stability #1104

Conversation

mcabbott commented Oct 16, 2021

simeonschaub commented Oct 16, 2021

mcabbott commented Oct 16, 2021

willtebbutt commented Oct 16, 2021

mcabbott commented Oct 16, 2021 • edited Loading

willtebbutt commented Oct 18, 2021

mcabbott left a comment • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

oxinabox commented Oct 26, 2021

mcabbott commented Oct 26, 2021 • edited Loading

oxinabox commented Oct 29, 2021

devmotion commented Oct 30, 2021

DhairyaLGandhi commented Oct 31, 2021

devmotion commented Oct 31, 2021

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

DhairyaLGandhi commented Nov 1, 2021 • edited Loading

devmotion commented Nov 1, 2021

Choose a reason for hiding this comment

mcabbott Nov 3, 2021 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

mcabbott Nov 3, 2021 • edited Loading

Choose a reason for hiding this comment

DhairyaLGandhi left a comment

Choose a reason for hiding this comment

ToucheSir commented Nov 7, 2021

mcabbott commented Nov 7, 2021

Insert `_project` into `getproperty`'s gradient, and then improve `z2d` etc. to restore stability #1104

Insert `_project` into `getproperty`'s gradient, and then improve `z2d` etc. to restore stability #1104

mcabbott commented Oct 16, 2021 •

edited

Loading

mcabbott left a comment •

edited

Loading

mcabbott commented Oct 26, 2021 •

edited

Loading

DhairyaLGandhi commented Nov 1, 2021 •

edited

Loading

mcabbott Nov 3, 2021 •

edited

Loading

mcabbott Nov 3, 2021 •

edited

Loading