Skip to content

Improvements to pullback_function #21

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Oct 11, 2021
Merged

Improvements to pullback_function #21

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
30 changes: 13 additions & 17 deletions src/AbstractDifferentiation.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -221,31 +221,27 @@ _zero(::AbstractVector, d::AbstractMatrix) = zero(similar(d, size(d, 2)))
_zero(::AbstractMatrix, d::AbstractMatrix) = zero(d)
_zero(::Any, d::Any) = zero(d)

@inline _dot(x, y) = dot(x, y)
@inline function _dot(x::AbstractVector, y::UniformScaling)
@assert length(x) == 1
return @inbounds dot(x[1], y.λ)
end
@inline function _dot(x::AbstractVector, y::AbstractMatrix)
@assert size(y, 2) == 1
return dot(x, y)
end

function pullback_function(ab::AbstractBackend, f, xs...)
return (ws) -> begin
jacs = jacobian(lowest(ab), (xs...,) -> begin
return gradient(lowest(ab), (xs...,) -> begin
vs = f(xs...)
if ws isa Tuple
@assert length(vs) == length(ws)
return sum(zip(vs, ws)) do v, w
if w isa Union{AbstractMatrix, UniformScaling} && v isa AbstractVector
return w' * v
else
# for arbitrary arrays
return dot(w, v)
end
end
return sum(Base.splat(_dot), zip(ws, vs))
else
w, v = ws, vs
if w isa Union{AbstractMatrix, UniformScaling} && v isa AbstractVector
return w' * v
else
# for arbitrary arrays
return dot(w, v)
end
return _dot(vs, ws)
end
end, xs...)
return adjoint.(jacs)
end
end
function value_and_pullback_function(
Expand Down