use Functors 0.3 in Flux (#2007)

* use Functors 0.3

* test which needs Functors 0.3

* fix show

* rm useless test

Author: mcabbott

Project.toml

@@ -29,7 +29,7 @@ Adapt = "3.0"
 ArrayInterface = "3.1, 4, 5, 6"
 CUDA = "3"
 ChainRulesCore = "1.12"
-Functors = "0.2.8"
+Functors = "0.3"
 MLUtils = "0.2"
 MacroTools = "0.5"
 NNlib = "0.8.9"

src/functor.jl

@@ -36,13 +36,16 @@ Possible values include:
 """
 trainmode!(m, mode = true) = mode isa Bool ? testmode!(m, !mode) : testmode!(m, mode)
 
-params!(p::Params, x::AbstractArray{<:Number}, seen = IdSet()) = push!(p, x)
-
 function params!(p::Params, x, seen = IdSet())
-  x in seen && return
-  push!(seen, x)
-  for child in trainable(x)
-    params!(p, child, seen)
+  if x isa AbstractArray{<:Number} && Functors.isleaf(x)
+    return push!(p, x)
+  elseif x in seen
+    nothing
+  else
+    push!(seen, x)
+    for child in trainable(x)
+      params!(p, child, seen)
+    end
   end
 end
 

src/layers/show.jl

@@ -47,7 +47,8 @@ end
 _show_leaflike(x) = isleaf(x)  # mostly follow Functors, except for:
 _show_leaflike(::Tuple{Vararg{<:Number}}) = true         # e.g. stride of Conv
 _show_leaflike(::Tuple{Vararg{<:AbstractArray}}) = true  # e.g. parameters of LSTMcell
-_show_leaflike(::Diagonal) = true                        # appears inside LayerNorm
+_show_leaflike(::Scale) = true                           # appears inside LayerNorm
+_show_leaflike(::AbstractArray{<:Number}) = true         # e.g. transposed arrays
 
 _show_children(x) = trainable(x)  # except for layers which hide their Tuple:
 _show_children(c::Chain) = c.layers

test/cuda/cuda.jl

@@ -112,7 +112,7 @@ end
 
   # Even more trivial: no movement
   @test gradient(x -> sum(abs, cpu(x)), a)[1] isa Matrix
-  @test gradient(x -> sum(abs, cpu(x)), a')[1] isa LinearAlgebra.Adjoint
+  @test gradient(x -> sum(abs, cpu(x)), a')[1] isa Matrix
   @test gradient(x -> sum(cpu(x)), a)[1] isa typeof(gradient(sum, a)[1]) # FillArray
   @test gradient(x -> sum(abs, gpu(x)), ca)[1] isa CuArray
   @test_skip gradient(x -> sum(abs, gpu(x)), ca')[1] isa CuArray # KernelError: passing and using non-bitstype argument

test/layers/show.jl

@@ -68,4 +68,8 @@ end
   @test Meta.isexpr(Meta.parse(toplevel_longchain), :call)  # comments are ignored
   @test Meta.parse(toplevel_longchain).args[1] == :Chain
 
+  # Functors@0.3 marks transposed matrices non-leaf, shouldn't affect printing:
+  adjoint_chain = repr("text/plain", Chain([Dense([1 2; 3 4]')]))
+  @test occursin("Dense(2 => 2)", adjoint_chain)
+  @test occursin("Chain([", adjoint_chain)
 end

test/utils.jl

@@ -665,6 +665,17 @@ end
   end
 end
 
+@testset "Shared parameters" begin
+  mat = [1 2; 3 4.0]
+  simple = ((nothing, mat, (3, mat, 4)))
+  @test length(Flux.params(simple)) == 1
+  
+  oneadj = (nt = (m = mat, a = mat'))
+  @test length(Flux.params(oneadj)) == 1  # needs Functors@0.3
+  
+  @test Flux.destructure(simple)[1] == Flux.destructure(oneadj)[1] == [1, 3, 2, 4]
+end
+
 @testset "Various destructure bugs" begin
 
   @testset "issue 1601" begin