WIP: Wirtinger support #54

simeonschaub · 2019-09-22T21:02:02Z

As discussed in #40, Wirtinger support is going to be moved out of master for now. I'm going to start working on it in this branch, but might eventually decide to move this into a package.

My current plans for this so far are:

remove type constraints for Wirtinger (since pretty much anything can be a differential, these don't make much sense to me anymore)
~~make wirtinger_[primal|conjugate] recursive, to work better for things like Thunks~~ add a function unthunk for this instead
introduce a type ComplexGradient, which works like Zygote's complex derivatives to address Special case derivative of non-holomorphic functions of type ℂ(^n)→ℝ #23 and make porting Zygote to ChainRules easier (Still needs docstrings)
introduce an abstract type AbstractWirtinger, which Wirtinger as well as ComplexGradient are a subtype of (Still needs docstrings)
a function chain, which works mostly like *, but respects which function is the derivative of the outer/inner function, which is important for AbstractWirtinger differentials (Still needs docstrings)
simplify @scalar_rule with this chain function
disallow multiplication of complex numbers with AbstractWirtinger and require users to use chain, since the order of chaining matters here, too (Still needs a better error)
tests
docs

I always appreciate any feedback.

This, together with adding `@inline` makes constant-propagation possible. Also fix a bug from before.

simeonschaub · 2019-09-27T14:13:07Z

Thinking about making ComplexGradient two real parameters instead of one complex one, so we could make use of Zero in some places. Would complicate some arithmetic, though.

nickrobinson251

I really like how this simplifies the code... but I admit i've not fully understood everything here yet. I think docs for ComplexGradient (and it's relation to Wirtinger) would be a big help.

nickrobinson251 · 2019-10-03T10:51:25Z

src/differential_arithmetic.jl

+    _chain(unthunk(outer), unthunk(inner), swap_order)
+
+@inline function _chain(outer, inner, swap_order)
+    if swap_order


why do we need to write things this way, rather than as e.g. if swap_order; a, b = b, a; end or swap_order && _chain(inner, outer, false)

The reason is that there's two orders here to consider. The first one, why we need the chain function at all, is concerning, which of the differentials is the partial of the outer and which one is the partial of the inner function. For AbstractWirtinger, this difference does matter, which is the purpose of chain. The second one is the order of multiplication, which matters if inner and outer are non-commutative objects like matrices. They might still be of type Wirtinger, only wirtinger_primal and wirtinger_conjugate are matrices. In general, both orderings are relevant, in Base.:* for example, we want to multiply the outer differential from the right, but this is not equivalent to chain(inner, outer). I will definitely explain this in a docstring.

src/differentials.jl

test/rules.jl

Co-Authored-By: Nick Robinson <[email protected]>

simeonschaub · 2019-10-03T12:07:49Z

I've written down some thoughts about this in FluxML/Zygote.jl#328 also. I still need to finish that though and will then try to adapt some of that in the docs here.

src/differentials.jl

Co-Authored-By: Nick Robinson <[email protected]>

oxinabox · 2019-10-18T14:36:35Z

I have not forgotten about this and started reading it again yesterday.

simeonschaub · 2019-10-19T15:07:28Z

test/differentials.jl


        # For most differentials, in most domains, this does nothing
-        for der in (DNE(), @thunk(23), @thunk(Wirtinger(2,2)), [1 2], One(), Zero(), 0.0)
+        for der in (DNE(), @thunk(23), [1 2], One(), Zero(), 0.0)


I don't know if this test really makes a lot of sense...

simeonschaub · 2019-10-19T15:10:54Z

src/differentials.jl

+All subtypes wrap real/holomorphic differentials, and should always be the outermost wrapper.
+E.g., a typical differential would look like this:
+```
+Wirtinger(@thunk(::AbstractArray{Number}), @thunk(::AbstractArray{<:Number}))


After thinking about this quite a bit, I think this order makes the most sense, since this way, we can avoid allocations completely for derivatives like Wirtinger(Zero(), One()).

simeonschaub · 2019-10-19T15:13:11Z

src/differentials.jl

+
+extern(x::AbstractWirtinger) = throw(ArgumentError("`AbstractWirtinger` cannot be converted to an external type."))
+
+Base.iterate(x::AbstractWirtinger) = (x, nothing)


I think, I am going to get rid of this, since e.g. Wirtinger can now wrap arrays as well.

I thought about making this behave like Wirtinger(i, j) for (i,j) in zip(primal, conjugate), but I think this would encourage people to collect this into an array, and we don't ever want to have AbstractArray{<:AbstractWirtinger}.

Yeah I think the iterate Overloads can probably be removed from everything

src/differentials.jl

src/rule_definition_tools.jl

oxinabox · 2019-10-19T20:11:04Z

src/differentials.jl

+
+extern(x::AbstractWirtinger) = throw(ArgumentError("`AbstractWirtinger` cannot be converted to an external type."))
+
+Base.iterate(x::AbstractWirtinger) = (x, nothing)


Yeah I think the iterate Overloads can probably be removed from everything

oxinabox · 2019-10-19T20:13:14Z

src/differentials.jl

+wirtinger_primal(::Union{AbstractThunk}) =
+    throw(ArgumentError("`wirtinger_primal` is not defined for `AbstractThunk`. Call `unthunk` first."))
+wirtinger_conjugate(::Union{AbstractThunk}) =
+    throw(ArgumentError("`wirtinger_conjugate` is not defined for `AbstractThunk`. Call `unthunk` first."))


Weird Union

Also not sure about this function but will wait and see

Yes, the Union is just a relic. Probably don't need this anymore, since Thunks should now never wrap AbstractWirtinger.

oxinabox · 2019-10-19T20:14:53Z

src/rule_definition_tools.jl

+Returns `@thunk body`, except for when `body` is a call to [`Wirtinger`](@ref) or [`ComplexGradient`](@ref).
+In this case, it is equivalent to `Wirtinger(@thunk(primal), @thunk(conjugate))` / `ComplexGradient(@thunk primal)`.
+"""
+macro _thunk(body)


This doesn't need to be a macro anymore as it is only called from with in a function of ASTs.

Could you take a look, whether everything is escaped correctly now? I'm not 100% sure, I know how escaping works.

Tests will break for now, and that's good.

YingboMa · 2020-01-07T04:02:45Z

disallow multiplication of complex numbers with AbstractWirtinger and require users to use chain, since the order of chaining matters here, too (Still needs a better error)

Why disallowing ::Wirtinger * ::Complex? How would a user interact with the Wirtinger object if there is no generic function defined?

simeonschaub · 2020-01-07T11:25:04Z

Most of the time, users shouldn't interact with Wirtinger objects at all. One of the main use case I see for them is as intermediary representations in mixed-mode AD implementations. If a library wants to expose that functionality to users, the authors can add their own abstraction on top to best fit their particular needs.
To understand, why we don't want to multiply complex numbers and Wirtinger objects, we can think of Wirtinger objects as 2x2 real Jacobians in a different basis, where two real numbers just happen to be represented by one complex number. In this context then, it is quite clear, that these form only a real vector space and why multiplication with a complex number is not well defined. We can still define an injective homomorphism from the complex numbers to Wirtinger objects by mapping z to Wirtinger(z, Zero()), but these don't commute with all Wirtinger objects anymore, which would be required for a complex vector space. JuliaDiff/ChainRules.jl#133 and JuliaDiff/ChainRules.jl#135 make it quite clear that subtypes of AbstractDifferential should form a vector space, so I believe disallowing ::Wirtinger * ::Complex is the only reasonable thing to do here.

willtebbutt · 2020-07-05T23:54:54Z

@simeonschaub what's the status of this? Has it been completely superceded by the other complex numbers stuff?

simeonschaub · 2020-07-06T06:50:08Z

What has been superseded is ComplexGradient, now that we just use Adjoint for that. It seems like people still expressed interest in Wirtinger derivatives, so this might be something to consider for v2. I currently don't have much use for this anymore, but if someone wanted to push this forward, I would certainly certainly offer to help wherever I can.

simeonschaub added 11 commits September 22, 2019 22:39

remove type constraints for Wirtinger

1548cbc

introduce AbstractWirtinger and ComplexGradient

88bb756

add chain function

e3ce538

make swap_order in chain a positional arg

186009a

This, together with adding `@inline` makes constant-propagation possible. Also fix a bug from before.

introduce a function unwrap_wirtinger

2618146

stop using types before they are defined

aa7a84a

fix tests

20e7134

rename unwrap_wirtinger -> unthunk

c42a953

fix chain function

2ac8801

use the new chain function in @scalar_rule

61a60ec

fix tests accordingly

41f1071

nickrobinson251 reviewed Oct 3, 2019

View reviewed changes

Update src/differentials.jl

93a7b14

Co-Authored-By: Nick Robinson <[email protected]>

add chain(::Real, ::ComplexGradient)

c904407

oxinabox reviewed Oct 6, 2019

View reviewed changes

src/differentials.jl Show resolved Hide resolved

oxinabox reviewed Oct 6, 2019

View reviewed changes

src/differentials.jl Show resolved Hide resolved

oxinabox reviewed Oct 6, 2019

View reviewed changes

src/differentials.jl Show resolved Hide resolved

Update src/differentials.jl

a83da2f

Co-Authored-By: Nick Robinson <[email protected]>

simeonschaub added 4 commits October 19, 2019 15:58

special case AbstractWirtinger in at_thunk

4d21e1a

add refine_differential for ComplexGradient

71e0c7b

overload Base.real for some differentials

491f781

add some docstrings

06ccb14

simeonschaub force-pushed the sim/wirtinger branch from b63f6ee to 06ccb14 Compare October 19, 2019 15:00

simeonschaub commented Oct 19, 2019

View reviewed changes

src/differentials.jl Outdated Show resolved Hide resolved

move at_thunk-magic into separate macro

6ce400c

oxinabox reviewed Oct 19, 2019

View reviewed changes

simeonschaub added 3 commits October 19, 2019 22:26

make at_scalar_rule detect wrong Wirtinger rules

b960a09

Tests will break for now, and that's good.

use _thunk as a function, not macro

3bc7e22

implement some of @oxinabox's suggestions

e3bf56d

simeonschaub force-pushed the sim/wirtinger branch from 8774ada to e3bf56d Compare October 19, 2019 21:25

simeonschaub mentioned this pull request Oct 20, 2019

remove comma in test_scalar JuliaDiff/ChainRules.jl#126

Merged

This was referenced Oct 29, 2019

Get rid of extern? / Add unthunk, use extern only for testing? #56

Closed

add Unthunk #64

Merged

This was referenced Jan 3, 2020

A trait-based system to handle Wirtinger derivatives #77

Open

Remove Wirtinger #79

Merged

nickrobinson251 added the Complex Differentiation Relating to any form of complex differentiation label Jun 17, 2020


		extern(x::AbstractWirtinger) = throw(ArgumentError("`AbstractWirtinger` cannot be converted to an external type."))

		Base.iterate(x::AbstractWirtinger) = (x, nothing)

WIP: Wirtinger support #54

Are you sure you want to change the base?

WIP: Wirtinger support #54

Uh oh!

Conversation

simeonschaub commented Sep 22, 2019 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

simeonschaub commented Sep 27, 2019

Uh oh!

nickrobinson251 left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

simeonschaub commented Oct 3, 2019

Uh oh!

Uh oh!

Uh oh!

Uh oh!

oxinabox commented Oct 18, 2019

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

YingboMa commented Jan 7, 2020

Uh oh!

simeonschaub commented Jan 7, 2020 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

willtebbutt commented Jul 5, 2020

Uh oh!

simeonschaub commented Jul 6, 2020

Uh oh!

Uh oh!

simeonschaub commented Sep 22, 2019 •

edited

Loading

simeonschaub commented Jan 7, 2020 •

edited

Loading