Patch Flux._isleaf for abstract arrays with bitstype elements (#2436)

* patch Flux._isleaf for Transpose and Adjoint arrays of bitstype elements

* patch `Flux._isleaf(::PermutedDimsArray)` as well

* `gradient(x -> sum(cpu(x)), ca')[1]` now returns `Adjoint{Float32, <:CuArray}` instead of plain `CuArray`

Author: jondeuce

src/functor.jl

@@ -181,12 +181,13 @@ julia> m.bias
 """
 cpu(x) = fmap(x -> adapt(FluxCPUAdaptor(), x), x, exclude = _isleaf)
 
-_isbitsarray(::AbstractArray{<:Number}) = true
-_isbitsarray(::AbstractArray{T}) where T = isbitstype(T)
-_isbitsarray(x) = false
+_isleaf(x) = Functors.isleaf(x)
+
+_isleaf(::AbstractArray{<:Number}) = true
+_isleaf(::AbstractArray{T}) where T = isbitstype(T)
+_isleaf(::Union{Transpose, Adjoint, PermutedDimsArray}) = false
 
 _isleaf(::AbstractRNG) = true
-_isleaf(x) = _isbitsarray(x) || Functors.isleaf(x)
 
 # the order below is important
 const GPU_BACKENDS = ("CUDA", "AMDGPU", "Metal", "CPU")

test/ext_cuda/cuda.jl

@@ -109,7 +109,7 @@ end
   # This test should really not go through indirections and pull out Fills for efficiency
   # but we forcefully materialise. TODO: remove materialising CuArray here
   @test gradient(x -> sum(cpu(x)), ca)[1] isa CuArray # This involves FillArray, which should be GPU compatible
-  @test gradient(x -> sum(cpu(x)), ca')[1] isa CuArray
+  @test gradient(x -> sum(cpu(x)), ca')[1] isa Adjoint{Float32, <:CuArray}
 
   # Even more trivial: no movement
   @test gradient(x -> sum(abs, cpu(x)), a)[1] isa Matrix

test/utils.jl

@@ -567,6 +567,30 @@ end
   @test length(Flux.params(oneadj)) == 1  # needs Functors@0.3
 
   @test Flux.destructure(simple)[1] == Flux.destructure(oneadj)[1] == [1, 3, 2, 4]
+
+  @testset "issue 2432" begin
+    x = rand(1)
+    m = (; a = x, b = x')
+    count = Ref(0)
+    mcopy = fmap(m; exclude = Flux._isleaf) do x
+      count[] += 1
+      return copy(x)
+    end
+    @test count[] == 1
+    @test mcopy.a === mcopy.b'
+
+    struct BitsType
+      x::Int32
+      y::Float64
+    end
+
+    for x in [1.0, 'a', BitsType(1, 2.0)]
+      @test Flux._isleaf([x])
+      @test !Flux._isleaf([x]')
+      @test !Flux._isleaf(transpose([x]))
+      @test !Flux._isleaf(PermutedDimsArray([x;;], (1, 2)))
+    end
+  end
 end
 
 @testset "Various destructure bugs" begin