Make fmap(f, x, y) useful

src/functor.jl

@@ -3,7 +3,7 @@ functor(T, x) = (), _ -> x
 functor(x) = functor(typeof(x), x)
 
 functor(::Type{<:Tuple}, x) = x, y -> y
-functor(::Type{<:NamedTuple}, x) = x, y -> y
+functor(::Type{<:NamedTuple{L}}, x) where L = NamedTuple{L}(map(s -> getproperty(x, s), L)), identity
 
 functor(::Type{<:AbstractArray}, x) = x, y -> y
 functor(::Type{<:AbstractArray{<:Number}}, x) = (), _ -> x
@@ -43,12 +43,11 @@ function _default_walk(f, x)
   re(map(f, func))
 end
 
-function fmap(f, x; exclude = isleaf, walk = _default_walk, cache = IdDict())
-  haskey(cache, x) && return cache[x]
-  y = exclude(x) ? f(x) : walk(x -> fmap(f, x, exclude = exclude, walk = walk, cache = cache), x)
-  cache[x] = y
+const INIT = Base._InitialValue()  # sentinel value for keyword not supplied
 
-  return y
+function fmap(f, x; exclude = isleaf, walk = _default_walk, cache = IdDict(), prune = INIT)
+  haskey(cache, x) && return prune === INIT ? cache[x] : prune
+  cache[x] = exclude(x) ? f(x) : walk(x -> fmap(f, x; exclude, walk, cache, prune), x)
 end
 
 ###
@@ -74,27 +73,16 @@ end
 ### Vararg forms
 ###
 
-function fmap(f, x, dx...; cache = IdDict())
-  haskey(cache, x) && return cache[x]
-  cache[x] = isleaf(x) ? f(x, dx...) : _default_walk((x...) -> fmap(f, x..., cache = cache), x, dx...)
+function fmap(f, x, ys...; exclude = isleaf, walk = _default_walk, cache = IdDict(), prune = INIT)
+  haskey(cache, x) && return prune === INIT ? cache[x] : prune
+  cache[x] = exclude(x) ? f(x, ys...) : walk((xy...,) -> fmap(f, xy...; exclude, walk, cache, prune), x, ys...)
 end
 
-function functor_tuple(f, x::Tuple, dx::Tuple)
-  map(x, dx) do x, x̄
-    _default_walk(f, x, x̄)
-  end
-end
-functor_tuple(f, x, dx) = f(x, dx)
-functor_tuple(f, x, ::Nothing) = x
-
-function _default_walk(f, x, dx)
+function _default_walk(f, x, ys...)
   func, re = functor(x)
-  map(func, dx) do x, x̄
-    # functor_tuple(f, x, x̄)
-    f(x, x̄)
-  end |> re
+  yfuncs = map(y -> functor(typeof(x), y)[1], ys)
+  re(map(f, func, yfuncs...))
 end
-_default_walk(f, ::Nothing, ::Nothing) = nothing
 
 ###
 ### FlexibleFunctors.jl
@@ -112,9 +100,7 @@ function makeflexiblefunctor(m::Module, T, pfield)
       func = NamedTuple{pfields}(map(p -> getproperty(x, p), pfields))
       return func, re
     end
-
   end
-
 end
 
 function flexiblefunctorm(T, pfield = :params)

test/basics.jl

@@ -1,25 +1,14 @@
-struct Foo
-  x
-  y
-end
+
+struct Foo; x; y; end
 @functor Foo
 
-struct Bar
-  x
-end
+struct Bar; x; end
 @functor Bar
 
-struct Baz
-  x
-  y
-  z
-end
-@functor Baz (y,)
+struct OneChild3; x; y; z; end
+@functor OneChild3 (y,)
 
-struct NoChildren 
-  x
-  y
-end
+struct NoChildren2; x; y; end
 
 @static if VERSION >= v"1.6"
   @testset "ComposedFunction" begin
@@ -31,6 +20,10 @@ end
   end
 end
 
+###
+### Basic functionality
+###
+
 @testset "Nested" begin
   model = Bar(Foo(1, [1, 2, 3]))
 
@@ -53,20 +46,48 @@ end
   @test fmap(f, x; exclude = x -> x isa AbstractArray) == x
 end
 
+@testset "Property list" begin
+  model = OneChild3(1, 2, 3)
+  model′ = fmap(x -> 2x, model)
+
+  @test (model′.x, model′.y, model′.z) == (1, 4, 3)
+end
+
+@testset "cache" begin
+  shared = [1,2,3]
+  m1 = Foo(shared, Foo([1,2,3], Foo(shared, [1,2,3])))
+  m1f = fmap(float, m1)
+  @test m1f.x === m1f.y.y.x
+  @test m1f.x !== m1f.y.x
+  m1p = fmapstructure(identity, m1; prune = nothing)
+  @test m1p == (x = [1, 2, 3], y = (x = [1, 2, 3], y = (x = nothing, y = [1, 2, 3])))
+
+  # A non-leaf node can also be repeated:
+  m2 = Foo(Foo(shared, 4), Foo(shared, 4))
+  @test m2.x === m2.y
+  m2f = fmap(float, m2)
+  @test m2f.x.x === m2f.y.x
+  m2p = fmapstructure(identity, m2; prune = Bar(0))
+  @test m2p == (x = (x = [1, 2, 3], y = 4), y = Bar(0))
+
+  # Repeated isbits types should not automatically be regarded as shared:
+  m3 = Foo(Foo(shared, 1:3), Foo(1:3, shared))
+  m3p = fmapstructure(identity, m3; prune = 0)
+  @test m3p.y.y == 0
+  @test_broken m3p.y.x == 1:3
+end
+
+###
+### Extras
+###
+
 @testset "Walk" begin
   model = Foo((0, Bar([1, 2, 3])), [4, 5])
 
   model′ = fmapstructure(identity, model)
   @test model′ == (; x=(0, (; x=[1, 2, 3])), y=[4, 5])
 end
 
-@testset "Property list" begin
-  model = Baz(1, 2, 3)
-  model′ = fmap(x -> 2x, model)
-
-  @test (model′.x, model′.y, model′.z) == (1, 4, 3)
-end
-
 @testset "fcollect" begin
   m1 = [1, 2, 3]
   m2 = 1
@@ -78,7 +99,7 @@ end
 
   m1 = [1, 2, 3]
   m2 = Bar(m1)
-  m0 = NoChildren(:a, :b)
+  m0 = NoChildren2(:a, :b)
   m3 = Foo(m2, m0)
   m4 = Bar(m3)
   @test all(fcollect(m4) .=== [m4, m3, m2, m1, m0])
@@ -89,6 +110,64 @@ end
   @test all(fcollect(m3) .=== [m3, m1, m2])
 end
 
+###
+### Vararg forms
+###
+
+@testset "fmap(f, x, y)" begin
+  m1 = (x = [1,2], y = 3)
+  n1 = (x = [4,5], y = 6)
+  @test fmap(+, m1, n1) == (x = [5, 7], y = 9)
+
+  # Reconstruction type comes from the first argument
+  foo1 = Foo([7,8], 9)
+  @test fmap(+, m1, foo1) == (x = [8, 10], y = 12)
+  @test fmap(+, foo1, n1) isa Foo
+  @test fmap(+, foo1, n1).x == [11, 13]
+
+  # Mismatched trees should be an error
+  m2 = (x = [1,2], y = (a = [3,4], b = 5))
+  n2 = (x = [6,7], y = 8)
+  @test_throws Exception fmap(first∘tuple, m2, n2)
+  @test_throws Exception fmap(first∘tuple, m2, n2)
+
+  # The cache uses IDs from the first argument
+  shared = [1,2,3]
+  m3 = (x = shared, y = [4,5,6], z = shared)
+  n3 = (x = shared, y = shared, z = [7,8,9])
+  @test fmap(+, m3, n3) == (x = [2, 4, 6], y = [5, 7, 9], z = [2, 4, 6])
+  z3 = fmap(+, m3, n3)
+  @test z3.x === z3.z
+end
+
+@testset "old test update.jl" begin
+  struct M{F,T,S}
+    σ::F
+    W::T
+    b::S
+  end
+
+  @functor M
+
+  (m::M)(x) = m.σ.(m.W * x .+ m.b)
+
+  m = M(identity, ones(Float32, 3, 4), zeros(Float32, 3))
+  x = ones(Float32, 4, 2)
+  m̄, _ = gradient((m,x) -> sum(m(x)), m, x)
+  m̂ = Functors.fmap(m, m̄) do x, y
+    isnothing(x) && return y
+    isnothing(y) && return x
+    x .- 0.1f0 .* y
+  end
+
+  @test m̂.W ≈ fill(0.8f0, size(m.W))
+  @test m̂.b ≈ fill(-0.2f0, size(m.b))
+end
+
+###
+### FlexibleFunctors.jl
+###
+
 struct FFoo
   x
   y
@@ -102,13 +181,13 @@ struct FBar
 end
 @flexiblefunctor FBar p
 
-struct FBaz
+struct FOneChild4
   x
   y
   z
   p
 end
-@flexiblefunctor FBaz p
+@flexiblefunctor FOneChild4 p
 
 @testset "Flexible Nested" begin
   model = FBar(FFoo(1, [1, 2, 3], (:y, )), (:x,))
@@ -132,7 +211,7 @@ end
 end
 
 @testset "Flexible Property list" begin
-  model = FBaz(1, 2, 3, (:x, :z))
+  model = FOneChild4(1, 2, 3, (:x, :z))
   model′ = fmap(x -> 2x, model)
 
   @test (model′.x, model′.y, model′.z) == (2, 2, 6)
@@ -147,7 +226,7 @@ end
   @test all(fcollect(m4, exclude = x -> x isa Array) .=== [m4, m3])
   @test all(fcollect(m4, exclude = x -> x isa FFoo) .=== [m4])
 
-  m0 = NoChildren(:a, :b)
+  m0 = NoChildren2(:a, :b)
   m1 = [1, 2, 3]
   m2 = FBar(m1, ())
   m3 = FFoo(m2, m0, (:x, :y,))

test/runtests.jl

@@ -5,7 +5,6 @@ using Zygote
 
   include("basics.jl")
   include("base.jl")
-  include("update.jl")
 
   if VERSION < v"1.6" # || VERSION > v"1.7-"
     @warn "skipping doctests, on Julia $VERSION"