From 96db615d76f343b48ebb74ba16d926d419d36bf5 Mon Sep 17 00:00:00 2001 From: Frames White Date: Tue, 16 Jan 2024 15:46:23 +0800 Subject: [PATCH] handle NoTangent in copyto! --- src/rulesets/Base/array.jl | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/src/rulesets/Base/array.jl b/src/rulesets/Base/array.jl index b94bddd3e..83f61bfca 100644 --- a/src/rulesets/Base/array.jl +++ b/src/rulesets/Base/array.jl @@ -86,7 +86,14 @@ _instantiate_zeros(ẋs::AbstractArray{<:AbstractArray}, xs) = ẋs ##### function frule((_, ẏ, ẋ), ::typeof(copyto!), y::AbstractArray, x) - return copyto!(y, x), copyto!(ẏ, ẋ) + if ẏ isa AbstractZero + # it's allowed to have an imutable zero tangent for ẏ as long as ẋ is zero + # TODO should this be handled here or in the AD? + @assert iszero(ẋ) + else + copyto!(ẏ, ẋ) + end + return copyto!(y, x), ẏ end function frule((_, ẏ, _, ẋ), ::typeof(copyto!), y::AbstractArray, i::Integer, x, js::Integer...)