derivative for matrix

Author: zhujch1

src/chainrules.jl

@@ -41,6 +41,17 @@ function rrule(::typeof(*), A::AbstractMatrix{S},
     return A * t, gemv_pullback
 end
 
+function rrule(::typeof(*), A::AbstractMatrix{S},
+    B::AbstractMatrix{TaylorScalar{T, N}}) where {N, S, T}
+    project_A = ProjectTo(A)
+    project_B = ProjectTo(B)
+    function gemm_pullback(x̄)
+        X̄ = unthunk(x̄)
+        NoTangent(), @thunk(project_A(X̄ * transpose(B))), @thunk(project_B(transpose(A) * X̄))
+    end
+    return A * B, gemm_pullback
+end
+
 @adjoint function +(t::Vector{TaylorScalar{T, N}}, v::Vector{T}) where {N, T}
     project_v = ProjectTo(v)
     t + v, x̄ -> (x̄, project_v(x̄))

src/derivative.jl

@@ -61,5 +61,6 @@ end
 
 @inline function derivative(f, x::AbstractMatrix{T}, l::AbstractVector{S},
     vN::Val{N}) where {T <: TN, S <: TN, N}
-    mapcols(u -> derivative(f, u, l, vN), x)
+    t = make_taylor.(x, l, vN)
+    return extract_derivative.(f(t), N)
 end