Functor Transpose et. al. (#33)

* simplest isbits usecache

* apply the cache only to leaf nodes

* functor Transpose, Adjoint, PermutedDimsArray

* rm version check, etc

* rm usecache, for another PR

* tidy, tests

Author: mcabbott

.github/workflows/ci.yml

@@ -17,7 +17,6 @@ jobs:
       fail-fast: false
       matrix:
         version:
-          - '1.0'
           - '1.6' # Replace this with the minimum Julia version that your package supports.
           - '1'   # automatically expands to the latest stable 1.x release of Julia
           - 'nightly'

Project.toml

@@ -1,11 +1,14 @@
 name = "Functors"
 uuid = "d9f16b24-f501-4c13-a1f2-28368ffc5196"
 authors = ["Mike J Innes <mike.j.innes@gmail.com>"]
-version = "0.2.8"
+version = "0.3.0"
+
+[deps]
+LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 
 [compat]
 Documenter = "0.27"
-julia = "1"
+julia = "1.6"
 
 [extras]
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"

src/Functors.jl

@@ -23,7 +23,6 @@ usually using the macro [`@functor`](@ref).
 """
 functor
 
-@static if VERSION >= v"1.5"  # var"@functor" doesn't work on 1.0, temporarily disable
 @doc """
     @functor T
     @functor T (x,)
@@ -66,7 +65,6 @@ TwoThirds(Foo(10, 20), Foo(3, 4), 560)
 ```
 """
 var"@functor"
-end  # VERSION
 
 """
     Functors.isleaf(x)

src/base.jl

@@ -1 +1,38 @@
+
 @functor Base.RefValue
+
+functor(::Type{<:Base.ComposedFunction}, x) = (outer = x.outer, inner = x.inner), y -> Base.ComposedFunction(y.outer, y.inner)
+
+###
+### Array wrappers
+###
+
+using LinearAlgebra
+# The reason for these is to let W and W' be seen as tied weights in Flux models.
+# Can't treat ReshapedArray very well, as its type doesn't include enough details for reconstruction.
+
+functor(::Type{<:Adjoint}, x) = (parent = _adjoint(x),), y -> adjoint(only(y))
+
+_adjoint(x) = adjoint(x)  # _adjoint is the inverse, and also understands more types:
+_adjoint(x::NamedTuple{(:parent,)}) = x.parent  # "structural" gradient, and lazy broadcast used by Optimisers:
+_adjoint(bc::Broadcast.Broadcasted{S}) where S = Broadcast.Broadcasted{S}(_conjugate(bc.f, adjoint), _adjoint.(bc.args))
+
+functor(::Type{<:Transpose}, x) = (parent = _transpose(x),), y -> transpose(only(y))
+
+_transpose(x) = transpose(x)
+_transpose(x::NamedTuple{(:parent,)}) = x.parent
+_transpose(bc::Broadcast.Broadcasted{S}) where S = Broadcast.Broadcasted{S}(_conjugate(bc.f, transpose), _transpose.(bc.args))
+
+_conjugate(f::F, ::typeof(identity)) where F = f
+_conjugate(f::F, op::Union{typeof(transpose), typeof(adjoint)}) where F = (xs...,) -> op(f(op.(xs)...))
+
+function functor(::Type{<:PermutedDimsArray{T,N,perm,iperm}}, x) where {T,N,perm,iperm}
+  (parent = _PermutedDimsArray(x, iperm),), y -> PermutedDimsArray(only(y), perm)
+end
+function functor(::Type{<:PermutedDimsArray{T,N,perm,iperm}}, x::PermutedDimsArray{Tx,N,perm,iperm}) where {T,Tx,N,perm,iperm}
+  (parent = parent(x),), y -> PermutedDimsArray(only(y), perm)  # most common case, avoid wrapping wrice.
+end
+
+_PermutedDimsArray(x, iperm) = PermutedDimsArray(x, iperm)
+_PermutedDimsArray(x::NamedTuple{(:parent,)}, iperm) = x.parent
+_PermutedDimsArray(bc::Broadcast.Broadcasted, iperm) = _PermutedDimsArray(Broadcast.materialize(bc), iperm)

src/functor.jl

@@ -8,10 +8,6 @@ functor(::Type{<:NamedTuple{L}}, x) where L = NamedTuple{L}(map(s -> getproperty
 functor(::Type{<:AbstractArray}, x) = x, y -> y
 functor(::Type{<:AbstractArray{<:Number}}, x) = (), _ -> x
 
-@static if VERSION >= v"1.6"
-  functor(::Type{<:Base.ComposedFunction}, x) = (outer = x.outer, inner = x.inner), y -> Base.ComposedFunction(y.outer, y.inner)
-end
-
 function makefunctor(m::Module, T, fs = fieldnames(T))
   yᵢ = 0
   escargs = map(fieldnames(T)) do f

test/base.jl

@@ -1,8 +1,86 @@
-@testset "Base" begin
-    @testset "RefValue" begin
-        x = Ref(1)
-        p, re = Functors.functor(x)
-        @test p == (x = 1,)
-        @test re(p) isa Base.RefValue{Int}
-    end
+
+@testset "RefValue" begin
+  @test fmap(sqrt, Ref(16))[] == 4.0
+  @test fmap(sqrt, Ref(16)) isa Ref
+  @test fmapstructure(sqrt, Ref(16)) === (x = 4.0,)
+
+  x = Ref(13)
+  p, re = Functors.functor(x)
+  @test p == (x = 13,)
+  @test re(p) isa Base.RefValue{Int}
+
+  x2 = (a = x, b = [7, x, nothing], c = (7, nothing, Ref(13)))
+  y2 = fmap(identity, x2)
+  @test x2.a !== y2.a  # it's a new Ref
+  @test y2.a === y2.b[2]  # relation is maintained
+  @test y2.a !== y2.c[3]  # no new relation created
+
+  x3 = Ref([3.14])
+  f3 = [Foo(x3, x), x3, x]
+  @test f3[1].x === f3[2]
+  y3 = fmapstructure(identity, f3)  # replaces mutable with immutable
+  @test y3[1].x === y3[2]
+  @test y3[1].x.x === y3[2].x
+  z3 = fmapstructure(identity, y3)
+  @test z3[1].x === z3[2]
+  @test z3[1].x.x === z3[2].x
+end
+
+@testset "ComposedFunction" begin
+  f1 = Foo(1.1, 2.2)
+  f2 = Bar(3.3)
+  @test Functors.functor(f1 ∘ f2)[1] == (outer = f1, inner = f2)
+  @test Functors.functor(f1 ∘ f2)[2]((outer = f1, inner = f2)) == f1 ∘ f2
+  @test fmap(x -> x + 10, f1 ∘ f2) == Foo(11.1, 12.2) ∘ Bar(13.3)
+end
+
+@testset "LinearAlgebra containers" begin
+  @test fmapstructure(identity, [1,2,3]') == (parent = [1, 2, 3],)
+  @test fmapstructure(identity, transpose([1,2,3])) == (parent = [1, 2, 3],)
+
+  CNT = Ref(0)
+  fv(x::Vector) = (CNT[]+=1; 10v)
+
+  v = [1,2,3]
+  nt = fmap(fv, (a=v, b=v', c=transpose(v), d=[1,2,3]'))
+
+  @test nt.a === adjoint(nt.b)  # does not break tie
+  @test nt.a === transpose(nt.c)
+
+  @test CNT[] == 2
+  @test nt.a == adjoint(nt.d)  # does not create a new tie
+  @test nt.a !== adjoint(nt.d)
+
+  @test nt.b isa Adjoint
+  @test nt.c isa Transpose
+
+  x = [1,2,3]'
+  xs = fmapstructure(identity, x)  # check it digests this, e.g. structural gradient representation
+  @test Functors.functor(typeof(x), xs) == Functors.functor(x)  # (no real need for [2] types to match)
+
+  x = transpose([1 2; 3 4])
+  yt = transpose([5 6; 7 8])
+  ym = Matrix(yt)  # check it digests this, e.g. simplest Matrix gradient
+  @test Functors.functor(typeof(x), yt)[1].parent == Functors.functor(typeof(x), ym)[1].parent
+
+  ybc = Broadcast.broadcasted(+, ym, 9)  # check it digests this, as Optimisers.jl makes these
+  collect(ybc) isa Vector
+  zbc = Functors.functor(typeof(x), ybc)[1].parent
+  @test zbc .+ 0 == Functors.functor(typeof(x), ym .+ 9)[1].parent
+
+  # Similar checks for Adjoint. 
+  x = adjoint([1 2im 3; 4im 5 6im])
+  yt = adjoint([7im 8 9; 0 im 2])
+  ym = Matrix(yt)
+  @test Functors.functor(typeof(x), yt)[1].parent == Functors.functor(typeof(x), ym)[1].parent
+
+  ybc = Broadcast.broadcasted(+, ym, [11im, 12, im])
+  collect(ybc) isa Vector
+  zbc = Functors.functor(typeof(x), ybc)[1].parent
+  @test zbc .+ 0 == Functors.functor(typeof(x), ym .+ [11im, 12, im])[1].parent
+end
+
+@testset "PermutedDimsArray" begin
+  @test fmapstructure(identity, PermutedDimsArray([1 2; 3 4], (2,1))) == (parent = [1 2; 3 4],)
+  @test fmap(exp, PermutedDimsArray([1 2; 3 4], (2,1))) isa PermutedDimsArray{Float64}
 end

test/basics.jl

@@ -4,23 +4,15 @@ using Functors: functor
 struct Foo; x; y; end
 @functor Foo
 
-struct Bar; x; end
+struct Bar{T}; x::T; end
 @functor Bar
 
 struct OneChild3; x; y; z; end
 @functor OneChild3 (y,)
 
 struct NoChildren2; x; y; end
 
-@static if VERSION >= v"1.6"
-  @testset "ComposedFunction" begin
-    f1 = Foo(1.1, 2.2)
-    f2 = Bar(3.3)
-    @test Functors.functor(f1 ∘ f2)[1] == (outer = f1, inner = f2)
-    @test Functors.functor(f1 ∘ f2)[2]((outer = f1, inner = f2)) == f1 ∘ f2
-    @test fmap(x -> x + 10, f1 ∘ f2) == Foo(11.1, 12.2) ∘ Bar(13.3)
-  end
-end
+struct NoChild{T}; x::T; end
 
 ###
 ### Basic functionality
@@ -187,7 +179,6 @@ end
   end
 end
 
-@static if VERSION >= v"1.6" # Julia 1.0: LoadError: error compiling top-level scope: type definition not allowed inside a local scope
 @testset "old test update.jl" begin
   struct M{F,T,S}
     σ::F
@@ -211,7 +202,6 @@ end
   @test m̂.W ≈ fill(0.8f0, size(m.W))
   @test m̂.b ≈ fill(-0.2f0, size(m.b))
 end
-end  # VERSION
 
 ###
 ### FlexibleFunctors.jl

test/runtests.jl

@@ -1,5 +1,6 @@
 using Functors, Test
 using Zygote
+using LinearAlgebra
 
 @testset "Functors.jl" begin