avoid allocating new array when reduce / repeat is a no-op

Author: AntonOresten

Project.toml

@@ -1,7 +1,7 @@
 name = "Einops"
 uuid = "e3ce28c8-8bfb-4704-8add-e3e7f14b55c9"
 authors = ["Anton Oresten <antonoresten@gmail.com> and contributors"]
-version = "0.1.5"
+version = "0.1.6"
 
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
@@ -14,4 +14,4 @@ ChainRulesCore = "1.7"
 EllipsisNotation = "1"
 OMEinsum = "0.8"
 TupleTools = "1"
-julia = "1.9"
+julia = "1.9"

src/reduce.jl

@@ -44,16 +44,16 @@ function Base.reduce(f::Function, x::AbstractArray, (left, right)::ArrowPattern;
     allunique(extract(Symbol, right)) || throw(ArgumentError("Right names $(right) are not unique"))
     left_names, right_names = extract(Symbol, left), extract(Symbol, right)
     expanded = reshape_in(x, left; context...)
-    reduced_dims::NTuple{length(left_names)-length(right_names),Int}, permutation = @ignore_derivatives begin
+    dims::NTuple{length(left_names)-length(right_names),Int}, permutation = @ignore_derivatives begin
         isempty(setdiff(right_names, left_names)) || throw(ArgumentError("All dimension names on right side of pattern must be present on left side: $(setdiff(right_names, left_names))"))
         reduced_dim_names = setdiff(left_names, right_names)
         reduced_dims = ntuple(i -> findfirst(isequal(reduced_dim_names[i]), left_names)::Int, length(left_names) - length(right_names))
         reduced_left_names = intersect(left_names, right_names)
         permutation = permutation_mapping(ntuple(i -> reduced_left_names[i], length(right_names)), right_names)
         reduced_dims, permutation
     end
-    reduced = f(expanded, dims=reduced_dims)
-    dropped = dropdims(reduced, dims=reduced_dims)
+    reduced = isempty(dims) ? expanded : f(expanded; dims)
+    dropped = dropdims(reduced; dims)
     permuted = _permutedims(dropped, permutation)
     collapsed = reshape_out(permuted, right)
     return collapsed

src/repeat.jl

@@ -51,7 +51,7 @@ function Base.repeat(x::AbstractArray, (left, right)::ArrowPattern; context...)
     expanded = reshape_in(x, left; context_info...)
     permuted = _permutedims(expanded, permutation)
     reshaped = reshape(permuted, prerepeat_shape(size(expanded), left_names, right_names))
-    repeated = repeat(reshaped, repeats...)
+    repeated = all(isone, repeats) ? reshaped : repeat(reshaped, repeats...)
     collapsed = reshape_out(repeated, right)
     return collapsed
 end

test/reduce.jl

@@ -139,4 +139,14 @@ using Test, Statistics
         # Can't drop dimensions in rearrange - use reduce instead
         @test reduce(sum, x, (:a, :b, :c, :d, :e) --> (:b, :d)) |> size == (100, 50)
     end
+
+    @testset "no-op allocation optimization" begin
+        x = rand(2, 3, 4)
+        
+        y = reduce(sum, x, (:a, :b, :c) --> (:a, :b, :c))
+        @test pointer(x) == pointer(y)
+
+        y = reduce(sum, x, (:a, :b, :c) --> ((:a, :b), :c))
+        @test pointer(x) == pointer(y)
+    end
 end

test/repeat.jl

@@ -1,5 +1,5 @@
 using Einops
-using Test
+using Test, Statistics
 
 @testset "Repeat Operations" begin
     @testset "basic repetitions" begin
@@ -91,4 +91,14 @@ using Test
         x = ["a" "b"; "c" "d"]
         @test repeat(x, (:a, :b) --> (:a, :b, :c), c=2) == cat(x, x, dims=3)
     end
+
+    @testset "no-op allocation optimization" begin
+        x = rand(2, 3, 4)
+
+        y = repeat(x, (:a, :b, :c) --> (:a, :b, :c, 1))
+        @test pointer(x) == pointer(y)
+
+        y = repeat(x, (:a, :b, :c) --> ((:a, :b), :c, 1))
+        @test pointer(x) == pointer(y)
+    end
 end