Flux.jl: einops in Flux.jl

[juliohm]

I am opening this issue to share the einops project with Flux.jl devs. We should definitely have something similar or more awesome in Julia. Are you aware of it? Do we have a similar solution in Julia already?

https://github.com/arogozhnikov/einops

[mcabbott]

I think that TensorSlice.jl is actually almost there. I added a macro @reduce (not yet in the readme), with which you can do things like the einops maxpool example:

using TensorSlice, Flux, ImageView, JuliennedArrays
imgs = Flux.Data.MNIST.images()[1:32]

@shape A[i\I, j\J] := imgs[(I,J)][i,j] J:8

@reduce B[i,j] := maximum(α,β) A[α\i, β\j] α:2, β:2, i:112÷2

imshow(A); imshow(B)

There are still a few rough edges but it basically works. Here [j\J]==[(j,J)] is a possible notation for reshaping, to help remember which is the "small" index. (Capital J is the grid dimension in A, and j the original pixels.) It should really be smart enough to infer the range of i.

In einops you can also introduce size=1 dimensions, which I think won't be hard to add here either -- you will just say B[i, 1, j, 1] := ....

I haven't tested a whole lot with Flux, but it appears to work, because it's just writing things like B = maximum(reshape(A,(2,56,2,112)), dims=(1,3)). It's possible that slicing operations will be awful with CuArrays.

Are there other things which einops does, which people find useful?

[juliohm]

I really like the string syntax that einops defines, it is much cleaner to read. Do you think you could achieve something similar with TensorSlice.jl maybe with a Julia macro and Symbols?

[mcabbott]

I do like how close @tensot / @einsum etc. are to the indexing expressions, but I confess that this @reduce idea is straying quite far from that. Putting sum(i,j) in the middle is meant to immitate \sum_{i,j} in latex.

Without actually parsing any strings, you could try to make a macro which would accept many arguments... something like this could be hooked up to my tensor_slice function, with some fiddling — do you fancy having a go?

macro redshape(exs...) @show(exs); nothing end  

Y = @redshape [X]  i j k l -n  ->  (i,k) (j,l) n   [i:2, j:3]

Z = @redshape [Y : sum, i:2, j:3]  i\k  j\l n  ->  k l n

Edit: better suggestion, I think you could extract what's needed unambiguously here, although it's far from what Julia thinks this means! I've added -n to mean reverse(..., dims=5). And I'm trying to make it understand ranges specified like sum(j:2, j:3) := ....

[juliohm]

Closing the issue since you are already tacking the problem in your project @mcabbott. Thanks!

[mcabbott]

Sure! Today it's registered as TensorCast.

Somewhere I have a notebook doing the einops tutorial which perhaps I should upload.