extend fmapstructure and add KeyPath, fmap_with_path, fmapstructure_with_path

[CarloLucibello]

Adding some utilities for working with paths along trees:

• KeyPath
• fmap_with_path
• fmapstructure_with_path

Since it came for free, I also extended fmapstructure to handle multiple inputs.

In Jax they have KeyPath and tree_map_with_path

TODO

• add tests for fmap_with_path
• add to docs

[darsnack]

This is great. We've talked about adding something like KeyPath to give nodes an "address" for things like weight sharing in Optimisers.jl (for immutable weights).

The KeyPath implementation look fine. But it would be nice to simplify the implementation so that we don't have two copies of every walk style.

Ultimately, the key path only matters for ExcludeWalk. What if we change the API so that every walk accepts a key path always? So, now the implementation for StructuralWalk is the body of StructuralWalkWithKeyPath, and we eliminate the latter. Finally, we have ExcludeWalk with a boolean flag indicating whether exclude accepts a key path or not. If false, then just drop passing in the key path. Alternatively, we have ExcludeWalk and ExcludeWalkWithoutPath.

It would be a breaking change, but I think it's cleaner and more maintainable in the long run.

[CarloLucibello]

But it would be nice to simplify the implementation so that we don't have two copies of every walk style.

So from the interface perspective, the suggestion would be to go from

fmap(x -> x.^2, model, exclude = x -> isleaf(x), walk = WalkNOTTakingKeyPath())
fmap_with_path((kp, x) -> x.^2, model, exclude = (kp, x) -> isleaf(x), walk = WalkTakingKeyPath())

to

fmap(x -> x.^2, model, exclude = x -> isleaf(x), walk = WalkTakingKeyPath())
fmap_with_path((kp, x) -> x.^2, model, exclude = (kp, x) -> isleaf(x), walk = WalkTakingKeyPath())

It seems reasonable. (maybe has an unnecessary performance hit on some exotic scenarios? I don't know).
Honestly though, I don't feel like venturing into tagging a breaking release soonish, without global considerations on what other breaking change we might put in. And I'd like to have these features in sooner rather than later.
Also, these walks take really few lines of code, so duplication doesn't seem too problematic.
So I would say to merge this as it is, and to start listing breaking changes proposals for next breaking release.

[darsnack]

Fair but @ToucheSir may want to weigh in on these changes and re: breakage now or later.

[ToucheSir]

My only comment for now would be whether separate _with_path functions were necessary, or if we could've gotten away with something like passing the keypath in as a keyword argument. We should really look into API cleanup because the increase in API surface area and code duplication is noticeable now. That however is a different issue/PR discussion.