Merge pull request #380 from sathvikbhagavan/sb/zygote_matrix

refactor: zygote rules for matrix inputs
SciML · Jan 21, 2025 · 301f600 · 301f600
2 parents c4de75e + 977f085
commit 301f600
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 2 deletions.
diff --git a/ext/DataInterpolationsChainRulesCoreExt.jl b/ext/DataInterpolationsChainRulesCoreExt.jl
@@ -64,9 +64,12 @@ function u_tangent(A::LinearInterpolation, t, Δ)
     out = zero.(A.u)
     idx = get_idx(A, t, A.iguesser)
     t_factor = (t - A.t[idx]) / (A.t[idx + 1] - A.t[idx])
-    if eltype(out) <: Number
+    if out isa AbstractVector{<:Number}
         out[idx] = Δ * (one(eltype(out)) - t_factor)
         out[idx + 1] = Δ * t_factor
+    elseif out isa AbstractMatrix
+        out[:, idx] = Δ * (one(eltype(out)) - t_factor)
+        out[:, idx + 1] = Δ * t_factor
     else
         @. out[idx] = Δ * (true - t_factor)
         @. out[idx + 1] = Δ * t_factor
@@ -91,10 +94,14 @@ function u_tangent(A::QuadraticInterpolation, t, Δ)
     Δt_rel₀ = t - A.t[i₀]
     Δt_rel₁ = t - A.t[i₁]
     Δt_rel₂ = t - A.t[i₂]
-    if eltype(out) <: Number
+    if out isa AbstractVector{<:Number}
         out[i₀] = Δ * Δt_rel₁ * Δt_rel₂ / (Δt₀ * Δt₂)
         out[i₁] = -Δ * Δt_rel₀ * Δt_rel₂ / (Δt₀ * Δt₁)
         out[i₂] = Δ * Δt_rel₀ * Δt_rel₁ / (Δt₂ * Δt₁)
+    elseif out isa AbstractMatrix
+        out[:, i₀] = Δ * Δt_rel₁ * Δt_rel₂ / (Δt₀ * Δt₂)
+        out[:, i₁] = -Δ * Δt_rel₀ * Δt_rel₂ / (Δt₀ * Δt₁)
+        out[:, i₂] = Δ * Δt_rel₀ * Δt_rel₁ / (Δt₂ * Δt₁)
     else
         @. out[i₀] = Δ * Δt_rel₁ * Δt_rel₂ / (Δt₀ * Δt₂)
         @. out[i₁] = -Δ * Δt_rel₀ * Δt_rel₂ / (Δt₀ * Δt₁)

diff --git a/test/zygote_tests.jl b/test/zygote_tests.jl
@@ -51,12 +51,18 @@ end
     t = collect(1.0:10.0)
     test_zygote(
         LinearInterpolation, u, t; name = "Linear Interpolation")
+    u2 = Matrix(hcat(u, u)')
+    test_zygote(
+        LinearInterpolation, u2, t; name = "Linear Interpolation with matrix input")
 end
 
 @testset "Quadratic Interpolation" begin
     u = [1.0, 4.0, 9.0, 16.0]
     t = [1.0, 2.0, 3.0, 4.0]
     test_zygote(QuadraticInterpolation, u, t; name = "Quadratic Interpolation")
+    u2 = Matrix(hcat(u, u)')
+    test_zygote(
+        QuadraticInterpolation, u2, t; name = "Quadratic Interpolation with matrix input")
 end
 
 @testset "Constant Interpolation" begin