Merge pull request #364 from sdaxen/sdaxen/real_imag_conj_array

Add adjoints for real, conj, and imag for arrays

src/lib/array.jl

@@ -196,6 +196,10 @@ end
   dropdims(xs, dims = dims), Δ -> (reshape(Δ, size(xs)...),)
 end
 
+@adjoint real(x::AbstractArray) = real(x), r̄ -> (real(r̄),)
+@adjoint conj(x::AbstractArray) = conj(x), r̄ -> (conj(r̄),)
+@adjoint imag(x::AbstractArray) = imag(x), ī -> (complex.(0, real.(ī)),)
+
 @adjoint function mean(xs::AbstractArray; dims = :)
   return mean(xs, dims=dims), Δ -> (_backmean(xs,Δ,dims),)
 end

test/gradcheck.jl

@@ -217,6 +217,21 @@ end
   @test gradtest(x -> dropdims(x, dims = (1, 2, 3)), rand(1, 1, 1, 3))
 end
 
+@testset "$f(::AbstractArray)" for f in (real, conj, imag)
+  rng, N = MersenneTwister(123456), 3
+  Ts = (Float64, ComplexF64)
+  @testset "$f(::Array{$IT})" for IT in Ts
+    A = randn(IT, N, N)
+    y, back = Zygote.pullback(f, A)
+    y2, back2 = Zygote.pullback(x->f.(x), A)
+    @test y == y2
+    @testset "back(::Array{$BT})" for BT in Ts
+      ȳ = randn(BT, N, N)
+      @test back(ȳ)[1] == back2(ȳ)[1]
+    end
+  end
+end
+
 @testset "(p)inv" begin
   rng, P, Q = MersenneTwister(123456), 13, 11
   A, B, C = randn(rng, P, Q), randn(rng, P, P), randn(Q, P)