fmapstructure

[ToucheSir]

Example usage:

m = Chain(Dense(10, 5), Dense(5, 2))

functorize(identity, m) |> typeof
# Tuple{NamedTuple{(:weight, :bias, :σ), Tuple{Matrix{Float32}, Vector{Float32}, typeof(identity)}}, NamedTuple{(:weight, :bias, :σ), Tuple{Matrix{Float32}, Vector{Float32}, typeof(identity)}}}

# Similar output to what you get from explicit gradients
functorize(x -> x isa AbstractArray ? x : nothing, m) |> typeof
# Tuple{NamedTuple{(:weight, :bias, :σ), Tuple{Matrix{Float32}, Vector{Float32}, Nothing}}, NamedTuple{(:weight, :bias, :σ), Tuple{Matrix{Float32}, Vector{Float32}, Nothing}}}

Note how the implementation is almost exactly the same as fmap, just without the re(construction). Ideally this functionality could be folded into fmap itself, but I'm not sure what the most elegant way to conditionally control reconstruct is.

cc @DhairyaLGandhi @darsnack

[DhairyaLGandhi]

Nice, this is basically just traversing the tree, right? We could technically replace the construction with an identity. maybe an indirection with fmap1? I wish we have better names for these functions.

Also, we would miss out on non-functor-able objects in this view which is a shame.

[ToucheSir]

Pretty much, I think we'd some kind of kwarg and an additional conditional for fmap, because re must be derived at runtime from x itself. Then there's the question of whether to pass that kwarg onto fmap1 and have it handle the conditional, or to make another function that works like fmap1 but ignores re.

Non-functorable objects are handled exactly the same way as they are in fmap, fcollect etc. Which is to say they're considered leaf nodes and retained in the tree as-is (assuming one uses exclude=isleaf).

[darsnack]

Is fmap1 exported? We could just change the name and 0% of users will notice. Could be called applyfandre? Or maybe something better than that.

I don't think complicating fmap is the way to go. Functors.jl has essentially three functions right now. I would be in favor of just adding a new function like factorize instead of more kwargs to fmap.

[darsnack]

If anything I feel like f here serves a similar role to exclude in fcollect. Would it make more sense to make fcollect maintain structure? You could easily flatten the nested tuples after the fact.

[ToucheSir]

I feel like f here serves a similar role to exclude in fcollect

Not really, exclude signals when to stop recursing while f is a transformation function. Hence fmap is really a structure preserving map. The reason why fmap and functorize have so much duplication is because they have different definitions of "structure preserving":

• fmap (currently) always maintains the original struct type. i.e. fmap(identity, m) == m.
• functorize maintains only the structure and values of the original struct.

Come to think of it, we could solve this just by having fmap1 as a parameter of fmap. Then the default function would be fmap1, and to replicate functorize you'd pass (f, x) -> map(f, children(x)). Now the question becomes what to call that parameter 😄

[darsnack]

Yeah but all the example transforms are just identity selectively applied. Well you convinced me anyways, because I think the API that you are proposing is like the "walk" parameter I asked for before. For example, here I could use this parameter to only walk the prune-able parameters. So, it's more flexible than just conditional reconstruction.

[ToucheSir]

fmap(length, m) (for parameter counting) is an example transform that isn't a selectively applied identity :)

walk isn't a bad name for the parameter. If we're all on the same page about this, I'll switch over the PR to adding that to fmap. We could still have a functorize(f, x; kwargs...) = fmap(f, x; walk=(f, x) -> map(f, children(x)), kwargs...) as well.

[darsnack]

Yeah switching the PR to walk sounds good. Maybe we can rename functorize to fmapstructure?

[darsnack]

Maybe add something in the README? Also, I think this can be lifted out of WIP; it's pretty much good to go 😄.

[darsnack]

Also should we rename fmap1 while we're at it? Could it be called walkfunctor? Or mapfunctor?

[ToucheSir]

I had to add a better test after finding a bug in the first version (walk wasn't being applied recursively). Also added some docstrings based on the fcollect docstring.

As for what to name fmap1, I'm honestly not sure. walkfunctor or mapfunctor seem like better names for fmap itself honestly. I thought of destructure_map_restructure, fmap_helper, default_fmap_helper, default_walk, etc. (all with and without leading underscores), but none of them really stood out.

[darsnack]

Well I came up with walkfunctor, because I plan to refactor FluxPrune.jl to have walkpruneable. Similarly, we could have walktrainable one day. walkleaves, walkchildren, mapleaves, and mapchildren are also options.

[darsnack]

Also, _fmap is still better than fmap1 despite being nearly the same.

[darsnack]

After sleeping on it, my vote is for default_walk.

[ToucheSir]

Thoughts on leading vs no leading underscore? I couldn't find any reference to fmap1 in the wild, but I imagine we will never want this to leak beyond the internals of Functors.jl.

[darsnack]

_default_walk and un-exported is probably best to keep it internal.

[ToucheSir]

Done. Shall we get this merged so that Flux can start using it?

[darsnack]

Left a couple docstring suggestions, but it looks good otherwise.

[darsnack]

Maybe we omit this? I'm not sure it's the right wording, since functored structs are still "plain-old Julia data structures."

[darsnack]

A couple of no-op line changes. Do you know what that's from?

[ToucheSir]

Not sure, my editor may have trimmed existing trailing whitespace automatically on save?

[darsnack]

Oh yeah, so we'll keep these

[darsnack]

Oh yeah, so we'll keep these

[darsnack]

Do we want a custom walk example?

[ToucheSir]

I couldn't think of a succinct one that wasn't fmapstructure, so I just added a doctest for that instead. Are there any examples you can think of?

[darsnack]

Maybe just something that only walks the children that are arrays of arrays? Or some other simple type constrained walk.

[ToucheSir]

Put in an absolutely uninspired example, but it's an example nonetheless. Happy to change it if something more illustrative comes up.

[darsnack]

@CarloLucibello I don't have permissions on this branch. Or do we use bors?

[DhairyaLGandhi]

Thanks all

[DhairyaLGandhi]

The name fmapstructure makes it sound like fmap itself doesn't preserve structure, which is incorrect. So a better name might be in order. An fmap is a pretty well established concept (see from Haskell)

[ToucheSir]

I agree the name isn't ideal, but the intent was to point out that this only preserves the structure (for some handwavy definition of "structure") and not the actual types involved. While fmap has signature Functor f => (a -> b) -> f a -> f b , this would be closer to Functor f => (a -> b) -> f a -> Either HList HMap (which is an invalid signature, but my point is to signal the change from a named type into a generic container).

Another way to look at this is as an instance of a poor man's hylomorphism. This take is appealing because it lets you view fcollect and destructure as catamorphisms (functor itself is basically unfold), and _restructure as an anamorphism (the re function returned from functor being equivalent to fold). The analogy is not perfect because of how different both languages' type systems are, but I think it's a good framework for thinking about the space of operations one could do with a (Functors.jl) functor.

[ToucheSir]

Actually, perhaps a more direct analogy would be that Functors.fmap with walk customization behaves more like traverse than fmap. In this framing, anything that implements functor would be not just a Functor, but a Traversable. fmapstructure would then be closer to calling traverse with a function that spits out (named)tuples (themselves Traversables).