Apply suggestions from code review

Co-authored-by: Brian Chen <[email protected]>
FluxML · Nov 3, 2022 · fec8950 · fec8950
1 parent 31cb97a
commit fec8950
Showing 1 changed file with 3 additions and 4 deletions.
diff --git a/src/maps.jl b/src/maps.jl
@@ -4,10 +4,9 @@ function fmap(f, x, ys...; exclude = isleaf,
                            walk = DefaultWalk(),
                            cache = IdDict(),
                            prune = NoKeyword())
-  _walk = if isnothing(cache)
-    ExcludeWalk(AnonymousWalk(walk), f, exclude)
-  else
-    CachedWalk(ExcludeWalk(AnonymousWalk(walk), f, exclude), prune, cache)
+  _walk = ExcludeWalk(AnonymousWalk(walk), f, exclude)
+  if isnothing(cache)
+    _walk = CachedWalk(_walk, prune, cache)
   end
   fmap(_walk, f, x, ys...)
 end