JuliennedArrays.jl
JuliennedArrays.jl copied to clipboard
add convenient functions `stack`/`unstack`
Unlike Align/Slices, the input to stack/unstack is the outer dimensions. It's possible to support multiple dimensions, but I didn't bother here because my direct target is one-dimensional case.
@DhairyaLGandhi this is a direct replacement for Flux.stack/Flux.unstack without allocating memories.
How does it behave with AD?
Good question. Without even trying it, I bet it doesn't work here. I was just thinking about the data preparation stage stack(imgs, 4) when I'm writing the image data layout blog for Flux.
This package is pretty lightweight so I believe @bramtayl won't want to have extra dependencies here to support AD.
Right, I tend to agree. I know that for LazyArrays, we had to support the apply! function which was the simple API that all of LazyArrays exposed. If JuliennedArrays has a similar format, we can add that somewhere too to keep the package light and AD ready.
To be clear, it might still work (LazyArrays did at the time) but the adjoint meant that it remained lazy in the backwards pass as well, so it was needed for efficiency.
I've opened another PR #28 for test updates and enhancements, thus 1b607d8 is unnecessary now.
Hi! I'm glad you find the package useful. I'd be very happy to accept a PR for auto-differentiation (although I'm not sure I could understand it enough to review it). I'm not sure I understand why stack/unstack are needed though?
Also, on the AD part, have you seen this? https://github.com/mcabbott/SliceMap.jl