`PartialFunction` Causing Issues with `Flux.gpu`

[JoshuaBillson]

It seems that PartialFunction is causing errors when moving the weights to the GPU with Flux.gpu.

Here's a MWE:

julia> using Flux, CUDA, Metalhead

julia> ResNet(18) |> Flux.gpu
ERROR: MethodError: no method matching PartialFunctions.PartialFunction(::String, ::typeof(Metalhead.addact), ::Tuple{typeof(relu)}, ::@NamedTuple{})
The type `PartialFunctions.PartialFunction` exists, but no method is defined for this combination of argument types when trying to construct it.

Closest candidates are:
  PartialFunctions.PartialFunction(::Any, ::Any, ::String, ::Any, ::T, ::N) where {T<:Tuple, N<:NamedTuple}
   @ PartialFunctions ~/.julia/packages/PartialFunctions/sxmEK/src/PartialFunctions.jl:22
  PartialFunctions.PartialFunction(::Any, ::T, ::N) where {T<:Tuple, N<:NamedTuple}
   @ PartialFunctions ~/.julia/packages/PartialFunctions/sxmEK/src/PartialFunctions.jl:18
  PartialFunctions.PartialFunction(::PartialFunctions.PartialFunction, ::Tuple, ::NamedTuple)
   @ PartialFunctions ~/.julia/packages/PartialFunctions/sxmEK/src/PartialFunctions.jl:26

Stacktrace:
  [1] macro expansion
    @ ~/.julia/packages/ConstructionBase/lUKuV/src/functions.jl:0 [inlined]
  [2] (::ConstructionBase.FunctionConstructor{PartialFunctions.PartialFunction{…}})(::String, ::typeof(Metalhead.addact), ::Tuple{typeof(relu)}, ::@NamedTuple{})
    @ ConstructionBase ~/.julia/packages/ConstructionBase/lUKuV/src/functions.jl:8
  [3] (::Functors.var"#3#6"{…})(y::@NamedTuple{…})
    @ Functors ~/.julia/packages/Functors/LbNAu/src/functor.jl:22
  [4] (::Functors.DefaultWalk)(::Function, ::Function)
    @ Functors ~/.julia/packages/Functors/LbNAu/src/walks.jl:73
  [5] ExcludeWalk
    @ ~/.julia/packages/Functors/LbNAu/src/walks.jl:126 [inlined]
  [6] (::Functors.CachedWalk{Functors.ExcludeWalk{…}, Functors.NoKeyword, Functors.WalkCache{…}})(::Function, ::Function)
    @ Functors ~/.julia/packages/Functors/LbNAu/src/walks.jl:177
  [7] (::Functors.var"#recurse#27"{Functors.CachedWalk{Functors.ExcludeWalk{…}, Functors.NoKeyword, Functors.WalkCache{…}}})(xs::Function)
    @ Functors ~/.julia/packages/Functors/LbNAu/src/walks.jl:54
  [8] map
    @ ./tuple.jl:356 [inlined]
  [9] map(::Function, ::@NamedTuple{connection::PartialFunctions.PartialFunction{…}, layers::Tuple{…}})
    @ Base ./namedtuple.jl:266
 [10] _map(::Function, ::@NamedTuple{connection::PartialFunctions.PartialFunction{…}, layers::Tuple{…}})
    @ Functors ~/.julia/packages/Functors/LbNAu/src/walks.jl:3
 [11] DefaultWalk
    @ ~/.julia/packages/Functors/LbNAu/src/walks.jl:73 [inlined]
 [12] ExcludeWalk
    @ ~/.julia/packages/Functors/LbNAu/src/walks.jl:126 [inlined]
 [13] CachedWalk
    @ ~/.julia/packages/Functors/LbNAu/src/walks.jl:177 [inlined]
 [14] (::Functors.var"#recurse#27"{Functors.CachedWalk{…}})(xs::Parallel{PartialFunctions.PartialFunction{…}, Tuple{…}})
    @ Functors ~/.julia/packages/Functors/LbNAu/src/walks.jl:54
 [15] map
    @ ./tuple.jl:356 [inlined]
 [16] _map(::Function, ::Tuple{Parallel{PartialFunctions.PartialFunction{…}, Tuple{…}}, Parallel{PartialFunctions.PartialFunction{…}, Tuple{…}}})
    @ Functors ~/.julia/packages/Functors/LbNAu/src/walks.jl:3
 [17] DefaultWalk
    @ ~/.julia/packages/Functors/LbNAu/src/walks.jl:73 [inlined]
 [18] ExcludeWalk
    @ ~/.julia/packages/Functors/LbNAu/src/walks.jl:126 [inlined]
 [19] CachedWalk
    @ ~/.julia/packages/Functors/LbNAu/src/walks.jl:177 [inlined]
 [20] recurse
    @ ~/.julia/packages/Functors/LbNAu/src/walks.jl:54 [inlined]
 [21] map
    @ ./tuple.jl:355 [inlined]
 [22] map
    @ ./namedtuple.jl:266 [inlined]
 [23] _map
    @ ~/.julia/packages/Functors/LbNAu/src/walks.jl:3 [inlined]
 [24] DefaultWalk
    @ ~/.julia/packages/Functors/LbNAu/src/walks.jl:73 [inlined]
 [25] ExcludeWalk
    @ ~/.julia/packages/Functors/LbNAu/src/walks.jl:126 [inlined]
 [26] CachedWalk
    @ ~/.julia/packages/Functors/LbNAu/src/walks.jl:177 [inlined]
 [27] recurse
    @ ~/.julia/packages/Functors/LbNAu/src/walks.jl:54 [inlined]
 [28] map
    @ ./tuple.jl:358 [inlined]
 [29] map(f::Functors.var"#recurse#27"{Functors.CachedWalk{…}}, t::Tuple{Chain{…}, Chain{…}, Chain{…}, Chain{…}, Chain{…}})
    @ Base ./tuple.jl:358
 [30] _map(::Function, ::Tuple{Chain{Tuple{…}}, Chain{Tuple{…}}, Chain{Tuple{…}}, Chain{Tuple{…}}, Chain{Tuple{…}}})
    @ Functors ~/.julia/packages/Functors/LbNAu/src/walks.jl:3
 [31] DefaultWalk
    @ ~/.julia/packages/Functors/LbNAu/src/walks.jl:73 [inlined]
 [32] ExcludeWalk
    @ ~/.julia/packages/Functors/LbNAu/src/walks.jl:126 [inlined]
 [33] CachedWalk
    @ ~/.julia/packages/Functors/LbNAu/src/walks.jl:177 [inlined]
 [34] recurse
    @ ~/.julia/packages/Functors/LbNAu/src/walks.jl:54 [inlined]
 [35] map
    @ ./tuple.jl:355 [inlined]
 [36] map
    @ ./namedtuple.jl:266 [inlined]
 [37] _map
    @ ~/.julia/packages/Functors/LbNAu/src/walks.jl:3 [inlined]
 [38] DefaultWalk
    @ ~/.julia/packages/Functors/LbNAu/src/walks.jl:73 [inlined]
 [39] ExcludeWalk
    @ ~/.julia/packages/Functors/LbNAu/src/walks.jl:126 [inlined]
 [40] CachedWalk
    @ ~/.julia/packages/Functors/LbNAu/src/walks.jl:177 [inlined]
 [41] (::Functors.var"#recurse#27"{Functors.CachedWalk{…}})(xs::Chain{Tuple{…}})
    @ Functors ~/.julia/packages/Functors/LbNAu/src/walks.jl:54
 [42] map
    @ ./tuple.jl:356 [inlined]
 [43] _map(::Function, ::Tuple{Chain{Tuple{…}}, Chain{Tuple{…}}})
    @ Functors ~/.julia/packages/Functors/LbNAu/src/walks.jl:3
 [44] DefaultWalk
    @ ~/.julia/packages/Functors/LbNAu/src/walks.jl:73 [inlined]
 [45] ExcludeWalk
    @ ~/.julia/packages/Functors/LbNAu/src/walks.jl:126 [inlined]
 [46] CachedWalk
    @ ~/.julia/packages/Functors/LbNAu/src/walks.jl:177 [inlined]
 [47] recurse
    @ ~/.julia/packages/Functors/LbNAu/src/walks.jl:54 [inlined]
 [48] map(f::Functors.var"#recurse#27"{Functors.CachedWalk{…}}, t::Tuple{Tuple{…}})
    @ Base ./tuple.jl:355
 [49] map
    @ ./namedtuple.jl:266 [inlined]
 [50] _map
    @ ~/.julia/packages/Functors/LbNAu/src/walks.jl:3 [inlined]
 [51] DefaultWalk
    @ ~/.julia/packages/Functors/LbNAu/src/walks.jl:73 [inlined]
 [52] ExcludeWalk
    @ ~/.julia/packages/Functors/LbNAu/src/walks.jl:126 [inlined]
 [53] CachedWalk
    @ ~/.julia/packages/Functors/LbNAu/src/walks.jl:177 [inlined]
 [54] (::Functors.var"#recurse#27"{Functors.CachedWalk{Functors.ExcludeWalk{…}, Functors.NoKeyword, Functors.WalkCache{…}}})(xs::Chain{Tuple{Chain{…}, Chain{…}}})
    @ Functors ~/.julia/packages/Functors/LbNAu/src/walks.jl:54
 [55] map
    @ ./tuple.jl:355 [inlined]
 [56] map(::Function, ::@NamedTuple{layers::Chain{Tuple{Chain{Tuple{…}}, Chain{Tuple{…}}}}})
    @ Base ./namedtuple.jl:266
 [57] _map(::Function, ::@NamedTuple{layers::Chain{Tuple{Chain{Tuple{…}}, Chain{Tuple{…}}}}})
    @ Functors ~/.julia/packages/Functors/LbNAu/src/walks.jl:3
 [58] (::Functors.DefaultWalk)(::Function, ::ResNet)
    @ Functors ~/.julia/packages/Functors/LbNAu/src/walks.jl:73
 [59] ExcludeWalk
    @ ~/.julia/packages/Functors/LbNAu/src/walks.jl:126 [inlined]
 [60] CachedWalk
    @ ~/.julia/packages/Functors/LbNAu/src/walks.jl:177 [inlined]
 [61] execute(::Functors.CachedWalk{Functors.ExcludeWalk{…}, Functors.NoKeyword, Functors.WalkCache{…}}, ::ResNet)
    @ Functors ~/.julia/packages/Functors/LbNAu/src/walks.jl:55
 [62] fmap(::Function, ::ResNet; exclude::Function, walk::Functors.DefaultWalk, cache::IdDict{Any, Any}, prune::Functors.NoKeyword)
    @ Functors ~/.julia/packages/Functors/LbNAu/src/maps.jl:11
 [63] fmap
    @ ~/.julia/packages/Functors/LbNAu/src/maps.jl:3 [inlined]
 [64] AbstractDevice
    @ ~/.julia/packages/MLDataDevices/dkOUm/src/public.jl:366 [inlined]
 [65] gpu
    @ ~/.julia/packages/Flux/Mhg1r/src/functor.jl:116 [inlined]
 [66] |>(x::ResNet, f::typeof(gpu))
    @ Base ./operators.jl:926
Some type information was truncated. Use `show(err)` to see complete types.

[ToucheSir]

Looking at the PartialFunction type, it appears one cannot be reconstructed from its fields. This means it won't work by default with ConstructionBase, which means it won't work by default with Functors. I think adding an extra constructor to PartialFunction would fix this, but I'm not sure. You may want to ask on their issue tracker if they're ok with this idea.

[JoshuaBillson]

I've opened a PR on PartialFunctions to solve this issue: archermarx/PartialFunctions.jl#12.

[JoshuaBillson]

The changes to PartialFunctions have been merged, and I've opened PR #290 to update the compatibility requirements in Project.toml.

[adrhill]

Sorry for somewhat derailing this issue, but I'm also running into issues with the implementation of addact using PartialFunctions. It's somewhat annoying to have to dispatch on the ugly type PartialFunctions.PartialFunction{nothing, nothing, typeof(Metalhead.addact), Tuple{typeof(relu)}, @NamedTuple{}}.

Why not adhere to the KISS principle and add an AddAct type?
Would also be great if all layer types for common pre-trained architectures lived in Flux.

[adrhill]

#287 is also related to the use of PartialFunction for addact.