Merge branch 'EnzymeAD:main' into main

Author: glou-nes

src/TracedRArray.jl

@@ -736,18 +736,18 @@ function Base.repeat(x::AnyTracedRArray{T,N}, counts::Vararg{Int,M}) where {T,N,
 
     # (d1, d2, ..., dP) -> (d1, 1, d2, 1, ..., dP, 1)
     interleaved_size = ones(Int, 2P)
-    interleaved_size[1:2:2N] .= size(x)
+    interleaved_size[1:2:(2N)] .= size(x)
 
     x_interleaved = reshape(x, interleaved_size...)
 
     # (d1, 1, d2, 1, ..., dP, 1) -> (d1, r1, d2, r2, ..., dP, rP)
     broadcast_target_size = interleaved_size
-    broadcast_target_size[2:2:2M] .= counts
+    broadcast_target_size[2:2:(2M)] .= counts
 
     x_broadcasted = broadcast_to_size(x_interleaved, broadcast_target_size)
 
     # (d1, r1, d2, r2, ..., dP, rP) -> (d1*r1, d2*r2, ..., dP*rP)
-    final_size = vec(prod(reshape(broadcast_target_size, 2, :), dims=1))
+    final_size = vec(prod(reshape(broadcast_target_size, 2, :); dims=1))
 
     x_final = reshape(x_broadcasted, final_size...)
 

test/basic.jl

@@ -366,8 +366,8 @@ end
 
 @testset "repeat" begin
     @testset for (size, counts) in Iterators.product(
-        [(2,), (2,3), (2,3,4), (2,3,4,5)],
-        [(), (1,), (2,), (2,1), (1,2), (2,2), (2,2,2), (1,1,1,1,1)]
+        [(2,), (2, 3), (2, 3, 4), (2, 3, 4, 5)],
+        [(), (1,), (2,), (2, 1), (1, 2), (2, 2), (2, 2, 2), (1, 1, 1, 1, 1)],
     )
         x = rand(size...)
         @test (@jit repeat(Reactant.to_rarray(x), counts...)) == repeat(x, counts...)