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

Jump to bottom

Remove incorrect push! and pop! gradients #1025

Merged
merged 9 commits into from
Sep 28, 2021
Merged

Remove incorrect push! and pop! gradients #1025

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
3d5fa62
fix push + pop gradient for vector of arrays, add real tests
mcabbott Jul 12, 2021
df7d069
tweak
mcabbott Jul 12, 2021
56b966f
allow only trivial gradients in push!(::Params) etc.
mcabbott Jul 12, 2021
dc1d15f
generalise, and fail
mcabbott Jul 12, 2021
4b3e448
fix
mcabbott Jul 12, 2021
182a30e
rm gradients which don't work
mcabbott Jul 18, 2021
1d1fbb2
rm unused methods from push(IdSet) gradient
mcabbott Jul 18, 2021
f906ef1
restrict push error to arrays, rm adjoint for params
mcabbott Jul 18, 2021
f9d17a1
Update test/features.jl
mcabbott Sep 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
14 changes: 0 additions & 14 deletions src/compiler/interface.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -183,20 +183,6 @@ function Base.push!(ps::Params, x)
return ps
end

@adjoint! function Base.push!(xs::IdSet, x...)
l = length(x)
push!(xs, x...), Δ -> begin
(Δ, ntuple(_ -> nothing, l)...)
end
end

@adjoint! function Base.push!(xs::Params, x::AbstractArray{T}...) where T
sz_x = size.(x)
push!(xs, x...), Δ -> begin
(Δ, map(x -> Ones{T}(x...), sz_x)...)
end
end

Base.push!(ps::Params, x...) = (foreach(x -> push!(ps, x), x); ps)

function Base.delete!(ps::Params, x)
Expand Down
23 changes: 2 additions & 21 deletions src/lib/array.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -74,27 +74,8 @@ _droplike(dy::Union{LinearAlgebra.Adjoint, LinearAlgebra.Transpose}, dxv::Abstra
_ -> error("Mutating arrays is not supported -- called copyto!(::$(typeof(xs)), _...)")

for f in [push!, pop!, pushfirst!, popfirst!]
@eval @adjoint! $f(xs, x...) = $f(xs, x...),
_ -> error("Mutating arrays is not supported -- called $($f)(::$(typeof(xs)), _...)")
end

# This is kind of bad, but at least we don't materialize the whole
# array. Prefer to use `Buffer`
# function _pullback(cx::Context, ::typeof(push!), xs::AbstractVector{<:AbstractArray}, x::AbstractArray{T}...) where T
@adjoint! function push!(xs::AbstractVector{<:AbstractArray}, x::AbstractArray{T}...) where T
sz_xs = size.(xs)
sz_x = size.(x)
push!(xs, x...), Δ -> begin
(Δ, map(x -> Ones{T}(x...), sz_x)...)
end
end

@adjoint! function pop!(xs::AbstractVector{<:AbstractArray{T}}) where T
sz_xs = size.(xs)
op = pop!(xs)
op, Δ -> begin
([Ones{T}(sz...) for sz in sz_xs], )
end
@eval @adjoint! $f(x::AbstractVector, ys...) = $f(x, ys...),
_ -> error("Mutating arrays is not supported -- called $($f)(::$(typeof(x)), _...)")
end

# General
Expand Down
16 changes: 0 additions & 16 deletions test/gradcheck.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1365,22 +1365,6 @@ using Zygote: Buffer
prod(copy(b))
end == (3,)

@testset "Limited Mutation" begin
p = [rand(3,3), rand(3,3)]
r = rand(5,5)

# TODO: ngradient cannot handle Vector{Array}
gs = gradient((p,x) -> sum(sum.(push!(p,x))), p, r)
@test length(p[end]) == length(gs[1][end])
@test gs[1] ≈ map(x -> one.(x), p)
@test gs[2] ≈ one.(r)

# p = [rand(3,3), rand(3,3)] # redefine `p` after mutation
# gs = gradient(x -> sum(pop!(x)), p)
# @test length(gs[1]) == 2
# @test gs[1][1] == one.(p[1])
end

end

@testset "FillArrays" begin
Expand Down