Merge pull request #225 from JuliaDiff/ox/adjoint

Add adjoint and transpose for thunks

Author: oxinabox

Project.toml

@@ -1,6 +1,6 @@
 name = "ChainRulesCore"
 uuid = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
-version = "0.9.10"
+version = "0.9.11"
 
 [deps]
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"

src/differentials/abstract_differential.jl

@@ -7,16 +7,6 @@ The subtypes of `AbstractDifferential` define a custom \"algebra\" for chain
 rule evaluation that attempts to factor various features like complex derivative
 support, broadcast fusion, zero-elision, etc. into nicely separated parts.
 
-All subtypes of `AbstractDifferential` implement the following operations:
-
-`+(a, b)`: linearly combine differential `a` and differential `b`
-
-`*(a, b)`: multiply the differential `b` by the scaling factor `a`
-
-`Base.conj(x)`: complex conjugate of the differential `x`
-
-`Base.zero(x) = Zero()`: a zero.
-
 In general a differential type is the type of a derivative of a value.
 The type of the value is for contrast called the primal type.
 Differential types correspond to primal types, although the relation is not one-to-one.
@@ -30,6 +20,18 @@ That allows for gradients to be accumulated.
 
 It generally also should be able to be added to a primal to give back another primal, as
 this facilitates gradient descent.
+
+All subtypes of `AbstractDifferential` implement the following operations:
+
+ - `+(a, b)`: linearly combine differential `a` and differential `b`
+ - `*(a, b)`: multiply the differential `b` by the scaling factor `a`
+ -  `Base.zero(x) = Zero()`: a zero.
+
+Further, they often implement other linear operators, such as `conj`, `adjoint`, `dot`.
+Pullbacks/pushforwards are linear operators, and their inputs are often
+`AbstractDifferential` subtypes.
+Pullbacks/pushforwards in-turn call other linear operators on those inputs.
+Thus it is desirable to have all common linear operators work on `AbstractDifferential`s.
 """
 abstract type AbstractDifferential end
 

src/differentials/thunks.jl

@@ -13,6 +13,34 @@ end
     return element, (val, new_state)
 end
 
+"""
+    @thunk expr
+
+Define a [`Thunk`](@ref) wrapping the `expr`, to lazily defer its evaluation.
+"""
+macro thunk(body)
+    # Basically `:(Thunk(() -> $(esc(body))))` but use the location where it is defined.
+    # so we get useful stack traces if it errors.
+    func = Expr(:->, Expr(:tuple), Expr(:block, __source__, body))
+    return :(Thunk($(esc(func))))
+end
+
+"""
+    unthunk(x)
+
+On `AbstractThunk`s this removes 1 layer of thunking.
+On any other type, it is the identity operation.
+
+In contrast to [`extern`](@ref) this is nonrecursive.
+"""
+@inline unthunk(x) = x
+
+@inline extern(x::AbstractThunk) = extern(unthunk(x))
+
+Base.conj(x::AbstractThunk) = @thunk(conj(unthunk(x)))
+Base.adjoint(x::AbstractThunk) = @thunk(adjoint(unthunk(x)))
+Base.transpose(x::AbstractThunk) = @thunk(transpose(unthunk(x)))
+
 #####
 ##### `Thunk`
 #####
@@ -59,42 +87,14 @@ Also if we did `Zero() * res[1]` then the result would be `Zero()` and `f(x)` wo
 with a field for each variable used in the expression, and call overloaded.
 
 Do not use `@thunk` if this would be equal or more work than actually evaluating the expression itself.
+This is commonly the case for scalar operators.
 
 For more details see the manual section [on using thunks effectively](http://www.juliadiff.org/ChainRulesCore.jl/dev/writing_good_rules.html#Use-Thunks-appropriately-1)
 """
 struct Thunk{F} <: AbstractThunk
     f::F
 end
 
-
-"""
-    @thunk expr
-
-Define a [`Thunk`](@ref) wrapping the `expr`, to lazily defer its evaluation.
-"""
-macro thunk(body)
-    # Basically `:(Thunk(() -> $(esc(body))))` but use the location where it is defined.
-    # so we get useful stack traces if it errors.
-    func = Expr(:->, Expr(:tuple), Expr(:block, __source__, body))
-    return :(Thunk($(esc(func))))
-end
-
-"""
-    unthunk(x)
-
-On `AbstractThunk`s this removes 1 layer of thunking.
-On any other type, it is the identity operation.
-
-In contrast to [`extern`](@ref) this is nonrecursive.
-"""
-@inline unthunk(x) = x
-
-@inline extern(x::AbstractThunk) = extern(unthunk(x))
-
-# have to define this here after `@thunk` and `Thunk` is defined
-Base.conj(x::AbstractThunk) = @thunk(conj(unthunk(x)))
-
-
 (x::Thunk)() = x.f()
 @inline unthunk(x::Thunk) = x()
 

test/differentials/composite.jl

@@ -99,7 +99,7 @@ end
 
         @testset "Tuples" begin
             @test ==(
-                typeof(Composite{Tuple{}}() + Composite{Tuple{}}()), 
+                typeof(Composite{Tuple{}}() + Composite{Tuple{}}()),
                 Composite{Tuple{}, Tuple{}}
             )
             @test (
@@ -219,7 +219,12 @@ end
 
     @testset "show" begin
         @test repr(Composite{Foo}(x=1,)) == "Composite{Foo}(x = 1,)"
-        @test repr(Composite{Tuple{Int,Int}}(1, 2)) == "Composite{Tuple{Int64,Int64}}(1, 2)"
+        # check for exact regex match not occurence( `^...$`) 
+        # and allowing optional whitespace (`\s?`)
+        @test occursin(
+            r"^Composite{Tuple{Int64,\s?Int64}}\(1,\s?2\)$",
+            repr(Composite{Tuple{Int64,Int64}}(1, 2)),
+        )
     end
 
     @testset "internals" begin

test/differentials/thunks.jl

@@ -36,6 +36,15 @@
         end
     end
 
+    @testset "Linear operators" begin
+        x_real = [2.0 4.0; 8.0 5.0]
+        x_complex = [(2.0 + im) 4.0; 8.0 (5.0 + 4im)]
+        @testset "$(typeof(x))" for x in (x_real, x_complex)
+            x_thunked = @thunk(1.0 * x)
+            @test unthunk(x_thunked') == x'
+            @test unthunk(transpose(x_thunked)) == transpose(x)
+        end
+    end
 
     @testset "Broadcast" begin
         @testset "Array" begin