handle nothing-nothing

Author: DhairyaLGandhi

src/functor.jl

@@ -46,7 +46,7 @@ children(x) = functor(x)[1]
 
 function functor_tuple(f, x::Tuple, dx::Tuple)
   map(x, dx) do x, x̄
-    fmap1(f, x, x̄)
+    _default_walk(f, x, x̄)
   end
 end
 functor_tuple(f, x, dx) = f(x, dx)
@@ -66,6 +66,7 @@ function _default_walk(f, x)
   func, re = functor(x)
   re(map(f, func))
 end
+_default_walk(f, ::Nothing, ::Nothing) = nothing
 
 """
     fmap(f, x; exclude = isleaf, walk = Functors._default_walk)
@@ -193,5 +194,5 @@ end
 # This way we can use `fmap` to update the params with their gradients
 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...)
+  cache[x] = isleaf(x) ? f(x, dx...) : _default_walk((x...) -> fmap(f, x..., cache = cache), x, dx...)
 end