Allow gradients in fmap

src/functor.jl

@@ -25,6 +25,24 @@ end
 
 isleaf(x) = functor(x)[1] === ()
 
+# for Chain
+function functor_tuple(f, x::Tuple, dx::Tuple)
+  map(x, dx) do x, x̄
+    fmap1(f, x, x̄)
+  end
+end
+functor_tuple(f, x, dx) = f(x, dx)
+functor_tuple(f, x, ::Nothing) = x
+
+# @functor Chain
+# Chain -> func = (layers = (Dense,Dense),), gs -> (layers...)
+function fmap1(f, x, dx)
+  func, re = functor(x)
+  map(func, dx) do x, x̄
+    functor_tuple(f, x, x̄)
+  end |> re
+end
+
 function fmap1(f, x)
   func, re = functor(x)
   re(map(f, func))
@@ -36,3 +54,8 @@ function fmap(f, x; cache = IdDict())
   haskey(cache, x) && return cache[x]
   cache[x] = isleaf(x) ? f(x) : fmap1(x -> fmap(f, x, cache = cache), x)
 end
+
+function fmap(f, x, dx...; cache = IdDict())
+  haskey(cache, x) && return cache[x]
+  cache[x] = isleaf(x) ? f(x, dx...) : fmap1((x...) -> fmap(f, x..., cache = cache), x, dx...)
+end