JuliennedArrays.jl icon indicating copy to clipboard operation
JuliennedArrays.jl copied to clipboard

add convenient functions `stack`/`unstack`

Open johnnychen94 opened this issue 4 years ago • 7 comments
trafficstars

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.

johnnychen94 avatar Aug 24 '21 14:08 johnnychen94

How does it behave with AD?

DhairyaLGandhi avatar Aug 24 '21 15:08 DhairyaLGandhi

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.

johnnychen94 avatar Aug 24 '21 15:08 johnnychen94

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.

DhairyaLGandhi avatar Aug 24 '21 15:08 DhairyaLGandhi

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.

DhairyaLGandhi avatar Aug 24 '21 15:08 DhairyaLGandhi

I've opened another PR #28 for test updates and enhancements, thus 1b607d8 is unnecessary now.

johnnychen94 avatar Aug 24 '21 16:08 johnnychen94

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?

bramtayl avatar Aug 25 '21 02:08 bramtayl

Also, on the AD part, have you seen this? https://github.com/mcabbott/SliceMap.jl

bramtayl avatar Aug 25 '21 03:08 bramtayl