Merge pull request #278 from DhairyaLGandhi/zyg

Add adjoint for `create_array`

Author: YingboMa

Project.toml

@@ -6,6 +6,7 @@ version = "0.11.0"
 [deps]
 AbstractAlgebra = "c3fe647b-3220-5bb0-a1ea-a7954cac585d"
 AbstractTrees = "1520ce14-60c1-5f80-bbc7-55ef81b5835c"
+ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 Combinatorics = "861a8166-3701-5b0c-9a16-15d98fcdc6aa"
 ConstructionBase = "187b0558-2788-49d3-abe0-74a17ed4e7c9"
 DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
@@ -22,6 +23,7 @@ TimerOutputs = "a759f4b9-e2f1-59dc-863e-4aeb61b1ea8f"
 [compat]
 AbstractAlgebra = "0.9, 0.10, 0.11, 0.12, 0.13, 0.14, 0.15"
 AbstractTrees = "0.3"
+ChainRulesCore = "0.9"
 Combinatorics = "1.0"
 ConstructionBase = "1.1"
 DataStructures = "0.18"
@@ -40,6 +42,7 @@ Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
 PkgBenchmark = "32113eaa-f34f-5b0d-bd6c-c81e245fc73d"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [targets]
-test = ["Test", "Random", "PkgBenchmark", "BenchmarkTools", "Pkg"]
+test = ["Test", "Random", "PkgBenchmark", "BenchmarkTools", "Pkg", "Zygote"]

src/SymbolicUtils.jl

@@ -53,4 +53,7 @@ include("api.jl")
 
 include("code.jl")
 
+# ADjoints
+include("adjoints.jl")
+
 end # module

src/adjoints.jl

@@ -0,0 +1,12 @@
+using ChainRulesCore
+import ChainRulesCore: rrule
+import .Code
+
+function rrule(::typeof(Code.create_array), A::Type{<:AbstractArray}, T, u::Val{j}, d::Val{dims}, elems...) where {dims, j}
+  y = Code.create_array(A, T, u, d, elems...)
+  function create_array_pullback(Δ)
+    dx = Δ
+    (NO_FIELDS, DoesNotExist(), DoesNotExist(), DoesNotExist(), DoesNotExist(), dx..., ntuple(_ -> DoesNotExist(), length(elems) - prod(dims) + j)...)
+  end
+  y, create_array_pullback
+end

test/adjoints.jl

@@ -0,0 +1,28 @@
+using SymbolicUtils, SymbolicUtils.Code
+using Zygote
+
+@testset "create_array adjoint" begin
+  elems = (1,2,3,4,5,)
+
+  Ts = (Float64, Float32, Float16, Int64, Int32)
+  dims_candidates = ((1, (2,3)), (2, (1,3)))
+  As = (Array,)
+
+  for T in Ts,
+      dims in dims_candidates,
+      A in As
+
+    u, dim = dims
+    ŷ, pb = Zygote.pullback(elems) do elems
+      SymbolicUtils.Code.create_array(A, T, Val(u), Val(dim), elems...)
+    end
+    y = SymbolicUtils.Code.create_array(A, T, Val(u), Val(dim), elems...)
+    @test y == ŷ
+
+    gs = pb(ones(T, length(elems)))
+    @test length(gs[1]) == length(elems)
+    for i = 1:(prod(dim)-1)
+      @test gs[1][i] == one(eltype(ŷ))
+    end
+  end
+end 

test/runtests.jl

@@ -21,3 +21,4 @@ include("code.jl")
 include("nf.jl")
 include("interface.jl")
 include("fuzz.jl")
+include("adjoints.jl")