LuxDL / Lux.jl

Elegant & Performant Scientific Machine Learning in Julia
https://lux.csail.mit.edu/
MIT License
486 stars 58 forks source link

Rewrite `@layer_map` to use KeyPath from Functors #615

Closed avik-pal closed 4 months ago

avik-pal commented 5 months ago

The current implementation for https://lux.csail.mit.edu/stable/api/Lux/contrib#Map-over-Layer was written when there was no such functionality in Functors. We can rewrite the @layer_map to use a custom Walk type and KeyPath.

Of course, as a transition, we will still provide a string to the function to avoid a breaking change. From v0.6, we will pass a keypath.

For anyone interested the relevant code is here https://github.com/LuxDL/Lux.jl/blob/main/src/contrib/map.jl