Constructor for collect(iterable layers)

[darsnack]

I came across this use case when working on FluxPrune.jl. If I have an array of functions that operate on layers, then I can do map((f, m) -> f(m), fs, model) when model is iterable like Chain or Parallel. But this returns a vector of layers instead of the type of model. For Chain, I can just pass the output of map back into the constructor, but this isn't the case for Parallel.

Of course, I can special case fs::AbstractVector and model::Chain in my case, but it would be nice to have a more generic way to do this. Basically something like re from Functors.jl but for the output of iterate(model).

[ToucheSir]

If I'm understanding the ask correctly, this should be doable with functor:

layers, re = functor(model)
transformed_layers = map((f, m) -> f(m), fs, layers)
model = re(transformed_layers)

[darsnack]

That's what I went to at first too, but the children from functor are not the same as the elements returned by iterable in general.

[ToucheSir]

That's certainly true, but I'm coming up blank for examples in the wild where it applies. For example, the only iterable layer in https://github.com/FluxML/Flux.jl/blob/master/src/layers/basic.jl is Chain. Parallel doesn't, and I do wonder if you'd want to @forward Parallel.layers Base.iterate because Parallel.connection may itself be a trainable layer.

[darsnack]

Parallel.connection may itself be a trainable layer.

True, so maybe iterating layers is the wrong abstraction. Maybe the right one is a collection of functions with the same structure as the model itself. The use-case I'm describing is how do you selectively apply functions to parts of a model in a generic way?

There's two parts I guess. First, a convenient way to write this structure of functions. Second, basically a generic version of update/apply from Optimisers.jl. Where you re urse both structures together and call each function on each node.

[darsnack]

This is the discussion that I think you are referencing. Though most of it was not recorded and on a ML call.

[ToucheSir]

FluxML/Functors.jl#1 handles the lockstep part, but doesn't help much with targeting a subset of parameters only.

I also found the thread in question, discussion starts from FluxML/Flux.jl#1509 (comment). There the problem statement was that {grad params} != {all params}, but you could easily see how that could be extended to {prunable params} != {grad params} != {all params} or any other way of defining a subset of parameters.

[darsnack]

Oh nice, that PR is basically what I was thinking when I mentioned update/apply from Optimisers.jl.

I think the new walk argument handles the discussion about parameter subsets. It seems like the combination of walk and the PR you linked might give us the foundation for what I am describing? You'd walk the functions and the parameters together. And your walk function is responsible for applying each function to each node?

[ToucheSir]

On further thought, I wonder if the pattern Flux uses for trainable and you use for pruneable isn't sufficient to get subsets working in conjunction with exclude/walk. That could be an alternative or even complementary solution for FluxML/Flux.jl#1293.

[darsnack]

Yeah I definitely think that's the case. We could even have a

function walksubset(f, x, subset)
    children, re = functor(x)
    subset_children = subset(x)
    re(map(x -> x ∈ subset_children ? f(x) : x, children))
end

eventually if this pattern becomes more common.