Merge pull request #1024 from mcabbott/paramcast

Make broadcasting over `Params` in the gradient an error

Author: mcabbott

src/compiler/interface.jl

@@ -201,6 +201,10 @@ end
 
 Base.Broadcast.broadcasted(f, ps::Params) = broadcasted(f, ps.order)
 
+@adjoint function Broadcast.broadcasted(f::Function, ps::Params)
+  f.(ps), _ -> throw(ArgumentError("Zygote.Params does not support broadcasting within gradients, try iteration `for p in ps`"))
+end
+
 Base.:(==)(x::Params, y::Params) = x.order.data == y.order.data
 
 function Base.show(io::IO, ps::Params)

src/lib/array.jl

@@ -73,14 +73,14 @@ _droplike(dy::Union{LinearAlgebra.Adjoint, LinearAlgebra.Transpose}, dxv::Abstra
 @adjoint getindex(::Type{T}, xs...) where {T} = T[xs...], dy -> (nothing, dy...)
 
 @adjoint! setindex!(xs::AbstractArray, x...) = setindex!(xs, x...),
-  _ -> error("Mutating arrays is not supported")
+  _ -> error("Mutating arrays is not supported -- called setindex!(::$(typeof(xs)), _...)")
 
 @adjoint! copyto!(args...) = copyto!(args...),
-  _ -> error("Mutating arrays is not supported")
+  _ -> error("Mutating arrays is not supported -- called copyto!(::$(typeof(xs)), _...)")
 
 for f in [push!, pop!, pushfirst!, popfirst!]
-  @eval @adjoint! $f(xs, x...) =
-    push!(xs, x...), _ -> error("Mutating arrays is not supported")
+  @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

test/interface.jl

@@ -38,7 +38,10 @@ using Zygote: Grads
     x, y = [1,2], [1]
     ps = Params([x, y])
     @test length.(ps) == length.([x, y]) # 617
+    @test size.(ps, 1) == [2, 1]
     @test all(Params([[1,1]]) .== Params([[1,1]]))
+
+    @test_throws ArgumentError gradient(() -> sum(sum.(ps)), ps)
   end
 
   @testset "indexing" begin