Update: I created an issue on Flux’s Github page and they solved the one dimensional case in a very elegant way, see here. To solve the multidimensional case one needs to write a fallback rule for getindex, see the linked issue. I will update this once I have a more concrete answer for the multidimensional case.
3 Likes