Skip to content
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

Fix maporbroadcast for ReverseDiff #142

Closed
wants to merge 5 commits into from
Closed

Fix maporbroadcast for ReverseDiff #142

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
88affba
fix maporbroadcast for ReverseDiff
mohamed82008 Oct 10, 2020
c238ce6
add test
mohamed82008 Oct 10, 2020
b53c8a4
fix Tracker + ReverseDiff case
mohamed82008 Oct 10, 2020
8be3f14
Update src/compat/reversediff.jl
mohamed82008 Oct 11, 2020
4bfdc81
Update test/runtests.jl
yebai Jan 28, 2021
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
2 changes: 0 additions & 2 deletions src/Bijectors.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -243,8 +243,6 @@ end

include("interface.jl")

# Broadcasting here breaks Tracker for some reason
maporbroadcast(f, x::AbstractArray{<:Any, N}...) where {N} = map(f, x...)
Comment on lines -246 to -247
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I still think it would make sense to keep this definition (and only change it for ReverseDiff). If the arrays are of the same dimension, it should be simpler to not invoke the broadcast machinery (for which adjoints require quite complex implementations e.g. in Tracker and Zygote that even make use of ForwardDiff).

Copy link
Member Author

@mohamed82008 mohamed82008 Oct 11, 2020 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Tracker's map also reduces to an array of TrackedReals even when the input is a TrackedArray. So broadcasting by default is needed for 2 out of 3 AD backends. So it makes sense as the default. Only when we have mixed TrackedArrays and arrays of TrackedReals that Tracker fails hence the need for the map workaround. ReverseDiff broadcasting doesn't fail in this case.

maporbroadcast(f, x::AbstractArray...) = f.(x...)

# optional dependencies
Expand Down
4 changes: 3 additions & 1 deletion src/compat/reversediff.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
module ReverseDiffCompat

using ..ReverseDiff: ReverseDiff, @grad, value, track, TrackedReal, TrackedVector,
TrackedMatrix
TrackedMatrix, TrackedArray
mohamed82008 marked this conversation as resolved.
Show resolved Hide resolved
using Requires, LinearAlgebra

using ..Bijectors: Log, SimplexBijector, maphcat, simplex_link_jacobian,
Expand Down Expand Up @@ -46,6 +46,8 @@ function Base.maximum(d::LocationScale{<:TrackedReal})
end
end

maporbroadcast(f, x::Union{AbstractArray, TrackedArray, AbstractArray{<:TrackedReal}}...) = f.(x...)

mohamed82008 marked this conversation as resolved.
Show resolved Hide resolved
logabsdetjac(b::Log{1}, x::Union{TrackedVector, TrackedMatrix}) = track(logabsdetjac, b, x)
@grad function logabsdetjac(b::Log{1}, x::AbstractVector)
return -sum(log, value(x)), Δ -> (nothing, -Δ ./ value(x))
Expand Down
2 changes: 2 additions & 0 deletions src/compat/tracker.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,8 @@ using .Tracker: Tracker,
using Compat: eachcol
using LinearAlgebra

# Broadcasting here breaks Tracker for some reason
maporbroadcast(f, x::Union{AbstractArray, TrackedArray, AbstractArray{<:TrackedReal}}...) = map(f, x...)
maporbroadcast(f, x::TrackedArray...) = f.(x...)
mohamed82008 marked this conversation as resolved.
Show resolved Hide resolved
mohamed82008 marked this conversation as resolved.
Show resolved Hide resolved
function maporbroadcast(
f,
Expand Down