move rule to Iterators.Zip to not break Zygote

JuliaDiff · Oct 19, 2022 · 51f1ed4 · 51f1ed4
1 parent 4f144fd
commit 51f1ed4
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/src/rulesets/Base/iterators.jl b/src/rulesets/Base/iterators.jl
@@ -21,7 +21,9 @@ end
 ##### `zip`
 #####
 
-function rrule(::typeof(zip), xs::AbstractArray...)
+# Attaching the rule to `zip` breaks Zygote, whose rule is on `Iterators.Zip`.
+# function rrule(::typeof(zip), xs::AbstractArray...)
+function rrule(::Type{<:Iterators.Zip}, xs::Tuple{Vararg{AbstractArray}})
     function zip_pullback(dy)
         @debug "zip array pullback" summary(dy)
         dxs = _tangent_unzip(unthunk(dy))
@@ -52,7 +54,7 @@ end
 
 # For testing
 function rrule(::ComposedFunction{typeof(collect), typeof(zip)}, xs::AbstractArray...)
-    y, back = rrule(zip, xs...)
+    y, back = rrule(Iterators.Zip, xs)
     return collect(y), back
 end