Add parse_shape

Author: AntonOresten

README.md

@@ -58,12 +58,13 @@ rgb_image = repeat(image, (:h, :w) --> (:repeat, :h, :w), repeat=3)
 
 *   [x] Implement `rearrange`.
 *   [x] Support Python implementation's string syntax for patterns with string macro.
+*   [x] Implement `parse_shape`.
+*   [x] Implement `pack` and `unpack`.
 *   [ ] Support ellipsis notation (using `..` from [EllipsisNotation.jl](https://github.com/SciML/EllipsisNotation.jl)).
 *   [ ] Implement `reduce`.
 *   [ ] Implement `repeat`.
 *   [ ] Explore integration with `PermutedDimsArray` or `TransmuteDims.jl` for lazy and statically inferrable permutations.
-*   [ ] Implement `einsum` (or wrap existing implementation) 
-*   [x] Implement `pack`, and `unpack`.
+*   [ ] Implement `einsum` (or wrap existing implementation).
 
 ## Contributing
 

src/Einops.jl

@@ -8,6 +8,7 @@ export ..
 
 include("utils.jl")
 export -->
+export parse_shape
 
 include("einops_str.jl")
 export @einops_str

src/einops_str.jl

@@ -1,80 +1,65 @@
 function parse_pattern(pattern::AbstractString)
-    occursin("->", pattern) || return tokenize_packing_pattern(pattern)
-
+    occursin("->", pattern) || return tokenize_generic(pattern)
     lhs, rhs = strip.(split(pattern, "->"; limit = 2))
+    occursin("->", rhs) && throw(ArgumentError("multiple \"->\" in pattern"))
     lhs_axes = tokenize_side(lhs)
     rhs_axes = tokenize_side(rhs)
     return Tuple(lhs_axes) --> Tuple(rhs_axes)
 end
 
 function tokenize_side(side::AbstractString)
+
+    function parse_token!(buf::IOBuffer, tokens::Vector)
+        if position(buf) > 0
+            s = String(take!(buf))
+            token = tryparse(Int, s)
+            isnothing(token) && (token = Symbol(s))
+            push!(tokens, token)
+        end
+    end
+
     tokens = Any[]
-    stack  = Vector{Any}[]
     buf    = IOBuffer()
+    stack  = Vector{Any}[]
     i = firstindex(side)
-
     while i <= lastindex(side)
         c = side[i]
-
         if c == ' '
-            if position(buf) > 0
-                s = String(take!(buf))
-                token = tryparse(Int, s)
-                isnothing(token) && (token = Symbol(s))
-                push!(tokens, token)
-            end
+            parse_token!(buf, tokens)
             i += 1
-
         elseif c == '('
-            if position(buf) > 0
-                s = String(take!(buf))
-                token = tryparse(Int, s)
-                isnothing(token) && (token = Symbol(s))
-                push!(tokens, token)
-            end
+            parse_token!(buf, tokens)
             push!(stack, tokens)
             tokens = Any[]
             i += 1
-
         elseif c == ')'
-            if position(buf) > 0
-                s = String(take!(buf))
-                token = tryparse(Int, s)
-                isnothing(token) && (token = Symbol(s))
-                push!(tokens, token)
-            end
+            parse_token!(buf, tokens)
             isempty(stack) && throw(ArgumentError("unmatched ')' in pattern"))
             sub = tokens
             tokens = pop!(stack)
             push!(tokens, Tuple(sub))
             i += 1
-
         elseif c == '.'
             # Expect literal "..."
             (i + 2 ≤ lastindex(side) && side[i:i+2] == "...") ||
                 throw(ArgumentError("single '.' not allowed in pattern"))
             push!(tokens, ..)
             i += 3
-
         else
             write(buf, c)
             i += 1
         end
     end
+    parse_token!(buf, tokens)
 
-    if position(buf) > 0
-        s = String(take!(buf))
-        token = tryparse(Int, s)
-        isnothing(token) && (token = Symbol(s))
-        push!(tokens, token)
-    end
     !isempty(stack) && throw(ArgumentError("unmatched '(' in pattern"))
     return Tuple(tokens)
 end
 
-function tokenize_packing_pattern(pattern::AbstractString)
-    Tuple(map(s -> s == "*" ? (*) : Symbol(s), filter(!isempty, split(pattern, ' '))))
-end
+const SpecialToken = Dict(:* => (*), :_ => (-))
+get_special_token(symbol) = get(SpecialToken, symbol, symbol)
+mapfilter(f, pred, xs) = map(f, filter(pred, xs))
+tokenize_generic(pattern) = Tuple(mapfilter(get_special_token ∘ Symbol, !isempty, split(pattern, ' ')))
 
 """
     einops"a (b c) -> (c b a)"

src/pack_unpack.jl

@@ -1,7 +1,12 @@
 const Wildcard = typeof(*)
 const PackPattern{N} = NTuple{N,Union{Symbol,Wildcard}}
 
-checkpacking(pattern::PackPattern) = count(x -> x isa Wildcard, pattern) == 1 || error("Only one wildcard (*) is allowed in the pattern")
+function check_packing_pattern(pattern::PackPattern)
+    count(x -> x isa Wildcard, pattern) == 1 || error("Only one wildcard (*) is allowed in the pattern")
+    allunique(pattern) || error("Pattern $(pattern) has duplicate elements")
+    return nothing
+end
+
 find_wildcard(pattern::PackPattern) = findfirst(x -> x isa Wildcard, pattern)
 
 size_before_wildcard(dims::Dims, pattern::PackPattern) = dims[1:find_wildcard(pattern)-1]
@@ -31,7 +36,7 @@ julia> packed_shapes
  (7, 9)
 """
 function pack(unpacked_arrays, pattern::PackPattern{N}) where N
-    checkpacking(pattern)
+    check_packing_pattern(pattern)
     reshaped_arrays = [reshape(A, packed_size(size(A), pattern)::Dims{N})::AbstractArray{<:Any,N} for A in unpacked_arrays]
     concatenated_array::AbstractArray{<:Any,N} = cat(reshaped_arrays..., dims=find_wildcard(pattern))
     packed_shapes = Dims[size_wildcard(size(unpacked_array), pattern) for unpacked_array in unpacked_arrays]
@@ -66,7 +71,7 @@ julia> unpack(packed_array, packed_shapes, (:i, :j, *)) .|> size
 ```
 """
 function unpack(packed_array::AbstractArray{<:Any,N}, packed_shapes, pattern::PackPattern{N}) where N
-    checkpacking(pattern)
+    check_packing_pattern(pattern)
     inds = Iterators.accumulate(+, Iterators.map(prod, packed_shapes))
     unpacked_arrays = map(Iterators.flatten((0, inds)), inds, packed_shapes) do i, j, ps
         unpacked_array = selectdim(packed_array, find_wildcard(pattern), i+1:j)

src/utils.jl

@@ -1,14 +1,13 @@
 struct Pattern{L,R} end
+
+# TODO: document -->
 (-->)(L, R) = Pattern{L, R}()
+
 Base.show(io::IO, ::Pattern{L,R}) where {L,R} = print(io, "$L --> $R")
 Base.iterate(::Pattern{L}) where L = (L, Val(:R))
 Base.iterate(::Pattern{<:Any,R}, ::Val{:R}) where R = (R, nothing)
 Base.iterate(::Pattern, ::Nothing) = nothing
 
-function permutation_mapping(left::NTuple{N,T}, right::NTuple{N,T}) where {N,T}
-    perm::Vector{Int} = findfirst.(isequal.([right...]), Ref([left...]))
-    return ntuple(i -> perm[i], Val(N))
-end
 
 extract(::Type, ::Tuple{}) = ()
 function extract(T::Type, input_tuple::Tuple)
@@ -24,6 +23,57 @@ function extract(T::Type, input_tuple::Tuple)
     return (instances_from_first..., extract(T, rest_elements)...)
 end
 
+
+@generated function findtype(::Type{T}, xs::Tuple) where T
+    inds = Int[]
+    for (i, el_type) in enumerate(xs.parameters)
+        el_type <: T && push!(inds, i)
+    end
+    return Expr(:tuple, inds...)
+end
+
+const Ignored = typeof(-)
+const ShapePattern{N} = NTuple{N,Union{Symbol,Ignored}}
+
+"""
+    parse_shape(x, pattern)
+
+Capture the shape of an array in a pattern by naming dimensions using `Symbol`s,
+and `-` to ignore dimensions.
+
+```jldoctest
+julia> parse_shape(rand(2,3,4), (:a, :b, -))
+(a = 2, b = 3)
+
+julia> parse_shape(rand(2,3), (-, -))
+NamedTuple()
+
+julia> parse_shape(rand(2,3,4,5), (:first, :second, :third, :fourth))
+(first = 2, second = 3, third = 4, fourth = 5)
+```
+
+The output is a `NamedTuple`, whose type contains the `Symbol` elements of the `pattern::NTuple{N,Union{Symbol,typeof(-)}}`,
+meaning that, unless the pattern is [constant-propagated](https://discourse.julialang.org/t/how-does-constant-propagation-work/22735/4),
+the output type is not known at compile time.
+
+`@code_warntype parse_shape(rand(2,3,4), (:a, :b, -))`
+
+`h() = parse_shape(rand(2,3,4), (:a, :b, -)); @code_warntype h()`
+"""
+function parse_shape(x::AbstractArray{<:Any,N}, pattern::ShapePattern{N}) where N
+    names = extract(Symbol, pattern)
+    allunique(names) || error("Pattern $(pattern) has duplicate elements")
+    inds = findtype(Symbol, pattern)
+    return NamedTuple{names,NTuple{length(inds),Int}}(size(x, i) for i in inds)
+end
+
+
+function permutation_mapping(left::NTuple{N,T}, right::NTuple{N,T}) where {N,T}
+    perm::Vector{Int} = findfirst.(isequal.([right...]), Ref([left...]))
+    return ntuple(i -> perm[i], Val(N))
+end
+
+
 # fix for 1.10:
 _permutedims(x::AbstractArray{T,0}, ::Tuple{}) where T = x
 _permutedims(x, perm) = permutedims(x, perm)

test/runtests.jl

@@ -15,27 +15,42 @@ using Test, Statistics
         @test repr((:a, :b, :c) --> (:c, :b, :a)) == "(:a, :b, :c) --> (:c, :b, :a)"
     end
 
+    @testset "parse_shape" begin
+        @test begin
+            x = rand(2,3,5)
+            shape = parse_shape(x, (:a, :b, :c))
+            (shape[:a], shape[:b], shape[:c]) == size(x)
+        end
+    end
+
     @testset "einops string tokenization" begin
 
+        @testset "parse_shape pattern" begin
+            @test einops"a _ c" == (:a, -, :c)
+            @test einops"_ _ _" == (-, -, -)
+        end
+
         @testset "arrow pattern" begin
             @test einops"a b c -> a (c b)" == ((:a, :b, :c) --> (:a, (:c, :b)))
             @test einops"a b c -> a(c b)" == ((:a, :b, :c) --> (:a, (:c, :b)))
             @test einops"a b 1 -> a 1 b" == ((:a, :b, 1) --> (:a, 1, :b))
             @test einops"a b () -> a () b" == ((:a, :b, ()) --> (:a, (), :b))
             @test einops"a b()->a()b" == ((:a, :b, ()) --> (:a, (), :b))
             @test einops"b ... -> b ..." == ((:b, ..) --> (:b, ..))
+            @test einops"b b -> a a" == ((:b, :b) --> (:a, :a))
             @test einops"->" == (() --> ())
             @test einops"-> 1" == (() --> (1,))
             @test_throws "'.'" Einops.parse_pattern("-> .")
             @test_throws "'('" Einops.parse_pattern("-> (")
             @test_throws "')'" Einops.parse_pattern("-> )")
         end
 
-        @testset "packing pattern" begin
+        @testset "pack and unpack pattern" begin
             @test einops"i j * k" == (:i, :j, *, :k)
             @test einops" i  j  *  k " == (:i, :j, *, :k)
             @test einops"* i" == (*, :i)
             @test einops"i *" == (:i, *)
+            @test einops"i i" == (:i, :i)
         end
 
     end