diff --git a/Project.toml b/Project.toml index deb03b09..c53e2c5e 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,6 @@ name = "DataInterpolations" uuid = "82cc6244-b520-54b8-b5a6-8a565e85f1d0" -version = "6.4.3" +version = "6.4.4" [deps] FindFirstFunctions = "64ca27bc-2ba2-4a57-88aa-44e436879224" diff --git a/ext/DataInterpolationsChainRulesCoreExt.jl b/ext/DataInterpolationsChainRulesCoreExt.jl index 9e9659d1..a811784d 100644 --- a/ext/DataInterpolationsChainRulesCoreExt.jl +++ b/ext/DataInterpolationsChainRulesCoreExt.jl @@ -64,8 +64,13 @@ 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]) - out[idx] = Δ * (one(eltype(out)) - t_factor) - out[idx + 1] = Δ * t_factor + if eltype(out) <: Number + out[idx] = Δ * (one(eltype(out)) - t_factor) + out[idx + 1] = Δ * t_factor + else + @. out[idx] = Δ * (true - t_factor) + @. out[idx + 1] = Δ * t_factor + end out end @@ -78,9 +83,15 @@ function u_tangent(A::QuadraticInterpolation, t, Δ) Δt₀ = t₁ - t₀ Δt₁ = t₂ - t₁ Δt₂ = t₂ - t₀ - out[i₀] = Δ * (t - A.t[i₁]) * (t - A.t[i₂]) / (Δt₀ * Δt₂) - out[i₁] = -Δ * (t - A.t[i₀]) * (t - A.t[i₂]) / (Δt₀ * Δt₁) - out[i₂] = Δ * (t - A.t[i₀]) * (t - A.t[i₁]) / (Δt₂ * Δt₁) + if eltype(out) <: Number + out[i₀] = Δ * (t - A.t[i₁]) * (t - A.t[i₂]) / (Δt₀ * Δt₂) + out[i₁] = -Δ * (t - A.t[i₀]) * (t - A.t[i₂]) / (Δt₀ * Δt₁) + out[i₂] = Δ * (t - A.t[i₀]) * (t - A.t[i₁]) / (Δt₂ * Δt₁) + else + @. out[i₀] = Δ * (t - A.t[i₁]) * (t - A.t[i₂]) / (Δt₀ * Δt₂) + @. out[i₁] = -Δ * (t - A.t[i₀]) * (t - A.t[i₂]) / (Δt₀ * Δt₁) + @. out[i₂] = Δ * (t - A.t[i₀]) * (t - A.t[i₁]) / (Δt₂ * Δt₁) + end out end @@ -100,7 +111,7 @@ function ChainRulesCore.rrule(::typeof(_interpolate), t::Number) deriv = derivative(A, t) function interpolate_pullback(Δ) - (NoTangent(), Tangent{typeof(A)}(; u = u_tangent(A, t, Δ)), deriv * Δ) + (NoTangent(), Tangent{typeof(A)}(; u = u_tangent(A, t, Δ)), sum(deriv .* Δ)) end return _interpolate(A, t), interpolate_pullback end