Fix rrule for matrix valued and multi-dimensional interpolation. (#487)

* Fix rrule for matrix valued and multi-dimensional interpolation. * Fix syntax for v1.3
JuliaMath · Apr 17, 2022 · a4e1c66 · a4e1c66 · mkitti · Apr 17, 2022
1 parent 3421847
commit a4e1c66
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 1 deletion.
diff --git a/src/chainrules/chainrules.jl b/src/chainrules/chainrules.jl
@@ -10,7 +10,7 @@ function ChainRulesCore.rrule(itp::AbstractInterpolation, x...)
     y = itp(x...)
     function interpolate_pullback(Δy)
         nope = ChainRulesCore.@not_implemented "`Interpolations.gradient` does not calculate a gradient with respect to the original data, only the evaluation point"
-        (nope, (Δy * Interpolations.gradient(itp, x...))...)
+        (nope, (sum(Δy .* g) for g in Interpolations.gradient(itp, x...))...)
     end
     y, interpolate_pullback
 end
diff --git a/test/chainrules.jl b/test/chainrules.jl
@@ -25,4 +25,15 @@ using Zygote
     else
         @test_skip Zygote.gradient(itp2,1,2) == Tuple(Interpolations.gradient(itp2,1,2))
     end
+
+    # Matrix-valued
+    y=[[[0 0; 0 0]] [[1 2; 0 0]];
+       [[0 0; 3 4]] [[1 2; 3 4]]]
+    itp_m = interpolate(y, BSpline(Linear()))
+    if VERSION ≥ v"1.3"
+        @test Zygote.jacobian(itp_m, 1,1) == Tuple((x->[x...]).(Interpolations.gradient(itp_m, 1,1)))
+    else
+        @test_skip Zygote.jacobian(itp_m, 1,1) == Tuple((x->[x...]).(Interpolations.gradient(itp_m, 1,1)))
+    end
+
 end