You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
function rrule(::typeof(/), A::AbstractVecOrMat{<:Real}, B::AbstractVecOrMat{<:Real})
# Compute C (return value differently) to just calling `A/B`
# so we have some parts precomputed
# that we need to use for computing the pullback
Aᵀ, dA_pb = rrule(adjoint, A)
Bᵀ, dB_pb = rrule(adjoint, B)
Cᵀ, dS_pb = rrule(\, Bᵀ, Aᵀ)
C, dC_pb = rrule(adjoint, Cᵀ)
function slash_pullback(Ȳ)
# These two run no matter which partial you want
# but they don't run if you only want the result (C)
_, dC = dC_pb(Ȳ)
_, dAᵀ, dBᵀ = dS_pb(dC)
# Each of these only run if you want that particular partial
∂A = @thunk last(dA_pb(dAᵀ))
∂B = @thunk last(dA_pb(dBᵀ))
(NO_FIELDS, ∂A, ∂B)
end
return C, slash_pullback
end
The text was updated successfully, but these errors were encountered:
With #30
this might be a great example
The text was updated successfully, but these errors were encountered: