support closures

Author: CarloLucibello

Project.toml

@@ -4,10 +4,12 @@ authors = ["Mike J Innes <mike.j.innes@gmail.com>"]
 version = "0.4.0"
 
 [deps]
+ConstructionBase = "187b0558-2788-49d3-abe0-74a17ed4e7c9"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 
 [compat]
+ConstructionBase = "1.4"
 Documenter = "0.27"
 julia = "1.6"
 

src/Functors.jl

@@ -1,5 +1,6 @@
 module Functors
 
+using ConstructionBase: constructorof
 export @functor, @flexiblefunctor, fmap, fmapstructure, fcollect
 
 include("functor.jl")

src/functor.jl

@@ -1,9 +1,9 @@
-function functor(T, x)
+function functor(T::Type, x)
   names = fieldnames(T)
   if isempty(names)
     return (), _ -> x
   end
-  S = T.name.wrapper # remove parameters from parametric types
+  S = constructorof(T) # remove parameters from parametric types and support anonymous functions
   vals = ntuple(i -> getfield(x, names[i]), length(names))
   return NamedTuple{names}(vals), y -> S(y...)
 end

test/basics.jl

@@ -148,6 +148,17 @@ end
   @test_throws Exception functor(NamedTuple{(:x, :y)}, (z=33, x=1))
 end
 
+@testset "anonymous functions" begin  
+  model = let W = rand(2,2), b = ones(2)
+    x -> tanh.(W*x .+ b)
+  end
+  newmodel = fmap(zero, model)
+  @test newmodel isa Function
+  @test newmodel([1,2]) == [0,0]
+  @test newmodel.W == [0 0; 0 0]
+  @test newmodel.b == [0, 0]  
+end
+
 ###
 ### Extras
 ###