Support Base.ComposedFunction

[oschulz]

Given

using Functors

struct Foo <: Function
    a::Float64
end

(f::Foo)(x) = f.a * x

@functor Foo

it would IMHO be very useful to support ComposedFunction, to enable

Functors.functor(Foo(4.2) ∘ Foo(5.3)) == ((outer = (a = 4.2,), inner = (a = 5.3,)), re)

If acceptable, I'd volunteer to implement it.

Update: Fixed example above

[DhairyaLGandhi]

We can compose arbitrary structs already, that's how Flux models work too. Should ComposedFunction be similar, or treated more like an invariant which the example seems to suggest.

[oschulz]

Sorry, my example above was a bit messed up, I fixed it now.

Currently, Functors.functor(Foo(4.2) ∘ Foo(5.3)) returns ((), var...), even though Functors.functor(Foo(4.2)) returns ((x = 4.2,), var...).

I would argue that we want Functors to be able to support standard Julia function composition, so that ComposedFunction can by used with Optimizers.jl, etc.

For example, with

using Zygote

ff = Foo(4.2) ∘ Foo(5.3)
Zygote.gradient((f, x) -> f(x), ff, 1.1)[1]

which yields

(outer = (a = 5.83,), inner = (a = 4.620000000000001,))

we should have a way back from the gradient (plus originial params) to a ComposedFunction{Foo,Foo}.

So Functors.functor(fff) should return ((outer = (a = 4.2,), inner = (a = 5.3,)), re), and re should accept a NamedTuple{(:outer, :inner)} and return a ComposedFunction{Foo,Foo}.

[darsnack]

Implementing this would just be writing @functor ComposedFunction somewhere.

What's more complicated is that so far, Functors.jl has not functor'd structs from elsewhere in the ecosystem. I suppose there is no downside to doing this...if you're using Functors.jl, then you want everything to be a functor when sensible.

[DhairyaLGandhi]

We can't functor most things. Especially the ones that have specialised methods that need to be invoked during optimisation. So things like Transpose, Adjoint, ReshapedArray, Dual, RGB (colorants), some structs like Atoms, AtomGraph etc

[darsnack]

Yeah but those are all examples of unconventional "leaf" nodes. ComposedFunction is probably not a leaf node. But I agree, it's safer not to functor in Functors.jl. If you need it to be functor'd, it is easy enough to add @functor ComposedFunction to your user code.

[oschulz]

We can't functor most things.

I agree, we could do it for ComposedFunction though, in a clean way, it wouldn't add to the load time of Functors in a measurable way and it's such a fundamental concept (layers of ML-models, etc.).

easy enough to add @functor ComposedFunction to your user code

That can't be done in packages though, it would be type piracy and if more than one packages does it ...

Also, I tried, @functor ComposedFunction isn't quite sufficient, but it would still be easy to implement.

[DhairyaLGandhi]

Yeah, as I mentioned it would need to look closer to a tuple or a chain.

[darsnack]

That can't be done in packages though, it would be type piracy and if more than one packages does it ...

I'm not sure how it would be better if Functors.jl did it for any struct. On the specific case of ComposedFunction, I think it is probably safe for us to do. Up to @DhairyaLGandhi if we add it or not.

I will say I don't see a huge benefit over Flux.Chain other than syntax.

Also, I tried, @functor ComposedFunction isn't quite sufficient, but it would still be easy to implement.

Ah, cause ComposedFunction has an inner constructor. Yeah, manually defining functor should be easy enough. Or add a NamedTuple constructor for ComposedFunction.

[oschulz]

I will say I don't see a huge benefit over Flux.Chain other than syntax.

Well, ComposedFunction is in Base, after all, and I think Functors.jl can potentially have wide applications beyond Flux.

Would a PR with

Functors.functor(::Type{<:ComposedFunction}, x) = (outer = x.outer, inner = x.inner), y -> y.outer ∘ y.inner 

(seems to work fine) be welcome, as a kind of Functors "chain-rule"?

After all, Functors.jl has explicit support for Tuple, NamedTuple, etc. as well, and function composition is a very fundamental concept.

[darsnack]

I'm not strongly opposed.

Slight bike shedding:

Functors.functor(::Type{<:ComposedFunction}, x) = (outer = x.outer, inner = x.inner), y -> ComposedFunction(y.outer, y.inner)

[oschulz]

Ah, yes, that's cleaner.

[DhairyaLGandhi]

I would be careful of such definitions, and we can't be liberal with such definitions because they can cause severe side effects. That said, composing functions is something that we handle and having a utility for Base.ComposedFunction should technically work. Hopefully it can handle the cases where inner and outer are themselves composed of a mixture of functored and non-functored entities. These would need tests, of course. We would also need to test much of the existing packages using Functors with this. I'm also trying to think if there is special handling of Base.ComposedFunction in Base or otherwise that we need to be aware of.

Yeah but those are all examples of unconventional "leaf" nodes

True, but these are just examples I had. It is entirely possible to come up with a definite case where its not leaf nodes that are involved. eltypes may not be a leaf node strictly because the implementation tags the array as the leaf instead. Most eltype definitions work either via method dispatches of the form f(x::AbstractArray{T <: SomeType}, ...) or via broadcasting. These still need to be dispatched to their respective methods. We can have specialisations on the Array level, eltype level and their respective broadcasting/IR counterparts.