Merge remote-tracking branch 'origin/master' into dg/grad

Author: DhairyaLGandhi

.github/workflows/CompatHelper.jl .github/workflows/CompatHelper.yml.github/workflows/CompatHelper.jl renamed to .github/workflows/CompatHelper.yml

@@ -1,7 +1,7 @@
 name = "Functors"
 uuid = "d9f16b24-f501-4c13-a1f2-28368ffc5196"
 authors = ["Mike J Innes <mike.j.innes@gmail.com>"]
-version = "0.2.0"
+version = "0.2.1"
 
 [deps]
 MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"

Project.toml

@@ -71,3 +71,22 @@ Any field not in the list will not be returned by `functor` and passed through a
 It is also possible to implement `functor` by hand when greater flexibility is required. See [here](https://github.com/FluxML/Functors.jl/issues/3) for an example.
 
 For a discussion regarding the need for a `cache` in the implementation of `fmap`, see [here](https://github.com/FluxML/Functors.jl/issues/2).
+
+Use `exclude` for more fine-grained control over whether `fmap` descends into a particular value (the default is `exclude = Functors.isleaf`):
+```julia
+julia> using CUDA
+
+julia> x = ['a', 'b', 'c'];
+
+julia> fmap(cu, x)
+3-element Array{Char,1}:
+ 'a': ASCII/Unicode U+0061 (category Ll: Letter, lowercase)
+ 'b': ASCII/Unicode U+0062 (category Ll: Letter, lowercase)
+ 'c': ASCII/Unicode U+0063 (category Ll: Letter, lowercase)
+
+julia> fmap(cu, x; exclude = x -> CUDA.isbitstype(eltype(x)))
+3-element CuArray{Char,1}:
+ 'a': ASCII/Unicode U+0061 (category Ll: Letter, lowercase)
+ 'b': ASCII/Unicode U+0062 (category Ll: Letter, lowercase)
+ 'c': ASCII/Unicode U+0063 (category Ll: Letter, lowercase)
+```

README.md

@@ -2,7 +2,7 @@ module Functors
 
 using MacroTools
 
-export @functor, fmap
+export @functor, fmap, fcollect
 
 include("functor.jl")
 

src/Functors.jl

@@ -29,7 +29,20 @@ macro functor(args...)
   functorm(args...)
 end
 
-isleaf(x) = functor(x)[1] === ()
+"""
+    isleaf(x)
+
+Return true if `x` has no [`children`](@ref) according to [`functor`](@ref).
+"""
+isleaf(x) = children(x) === ()
+
+"""
+    children(x)
+
+Return the children of `x` as defined by [`functor`](@ref).
+Equivalent to `functor(x)[1]`.
+"""
+children(x) = functor(x)[1]
 
 # for Chain
 function functor_tuple(f, x::Tuple, dx::Tuple)
@@ -56,9 +69,67 @@ end
 
 # See https://github.com/FluxML/Functors.jl/issues/2 for a discussion regarding the need for
 # cache.
-function fmap(f, x; cache = IdDict())
+function fmap(f, x; exclude = isleaf, cache = IdDict())
   haskey(cache, x) && return cache[x]
-  cache[x] = isleaf(x) ? f(x) : fmap1(x -> fmap(f, x, cache = cache), x)
+  y = exclude(x) ? f(x) : fmap1(x -> fmap(f, x, cache = cache, exclude = exclude), x)
+  cache[x] = y
+
+  return y
+end
+
+"""
+    fcollect(x; exclude = v -> false)
+
+Traverse `x` by recursing each child of `x` as defined by [`functor`](@ref)
+and collecting the results into a flat array.
+
+Doesn't recurse inside branches rooted at nodes `v`
+for which `exclude(v) == true`.
+In such cases, the root `v` is also excluded from the result.
+By default, `exclude` always yields `false`. 
+
+See also [`children`](@ref).
+
+# Examples
+
+```jldoctest
+julia> struct Foo; x; y; end
+
+julia> @functor Foo
+
+julia> struct Bar; x; end
+
+julia> @functor Bar
+
+julia> struct NoChildren; x; y; end 
+
+julia> m = Foo(Bar([1,2,3]), NoChildren(:a, :b))
+
+julia> fcollect(m)
+4-element Vector{Any}:
+ Foo(Bar([1, 2, 3]), NoChildren(:a, :b))
+ Bar([1, 2, 3])
+ [1, 2, 3]
+ NoChildren(:a, :b)
+
+julia> fcollect(m, exclude = v -> v isa Bar)
+2-element Vector{Any}:
+ Foo(Bar([1, 2, 3]), NoChildren(:a, :b))
+ NoChildren(:a, :b)
+ 
+julia> fcollect(m, exclude = v -> Functors.isleaf(v))
+2-element Vector{Any}:
+ Foo(Bar([1, 2, 3]), NoChildren(:a, :b))
+ Bar([1, 2, 3])
+```
+"""
+function fcollect(x; cache = [], exclude = v -> false)
+  x in cache && return cache
+  if !exclude(x)
+    push!(cache, x)
+    foreach(y -> fcollect(y; cache = cache, exclude = exclude), children(x))
+  end
+  return cache
 end
 
 # Allow gradients and other constructs that match the structure of the functor

src/functor.jl

@@ -1,19 +1,29 @@
 using Functors, Test
 
-@testset "Nested" begin
-  struct Foo
-    x
-    y
-  end
+struct Foo
+  x
+  y
+end
+@functor Foo
 
-  @functor Foo
+struct Bar
+  x
+end
+@functor Bar
 
-  struct Bar
-    x
-  end
+struct Baz
+  x
+  y
+  z
+end
+@functor Baz (y,)
 
-  @functor Bar
+struct NoChildren 
+  x
+  y
+end
 
+@testset "Nested" begin
   model = Bar(Foo(1, [1, 2, 3]))
 
   model′ = fmap(float, model)
@@ -22,17 +32,40 @@ using Functors, Test
   @test model′.x.y isa Vector{Float64}
 end
 
+@testset "Exclude" begin
+  f(x::AbstractArray) = x
+  f(x::Char) = 'z'
+
+  x = ['a', 'b', 'c']
+  @test fmap(f, x)  == ['z', 'z', 'z']
+  @test fmap(f, x; exclude = x -> x isa AbstractArray) == x
+
+  x = (['a', 'b', 'c'], ['d', 'e', 'f'])
+  @test fmap(f, x)  == (['z', 'z', 'z'], ['z', 'z', 'z'])
+  @test fmap(f, x; exclude = x -> x isa AbstractArray) == x
+end
+
 @testset "Property list" begin
-  struct Baz
-    x
-    y
-    z
-  end
-  
-  @functor Baz (y,)
-  
   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
+  m3 = Foo(m1, m2)
+  m4 = Bar(m3)
+  @test all(fcollect(m4) .=== [m4, m3, m1, m2])
+  @test all(fcollect(m4, exclude = x -> x isa Array) .=== [m4, m3, m2])
+  @test all(fcollect(m4, exclude = x -> x isa Foo) .=== [m4])
+
+  m1 = [1, 2, 3]
+  m2 = Bar(m1)
+  m0 = NoChildren(:a, :b)
+  m3 = Foo(m2, m0)
+  m4 = Bar(m3)
+  println(fcollect(m4))
+  @test all(fcollect(m4) .=== [m4, m3, m2, m1, m0])
+end