Skip to content

Commit b56e661

Browse files
chore: format code (EnzymeAD#371)
Co-authored-by: mofeing <[email protected]>
1 parent 311498b commit b56e661

File tree

2 files changed

+5
-5
lines changed

2 files changed

+5
-5
lines changed

src/TracedRArray.jl

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -736,18 +736,18 @@ function Base.repeat(x::AnyTracedRArray{T,N}, counts::Vararg{Int,M}) where {T,N,
736736

737737
# (d1, d2, ..., dP) -> (d1, 1, d2, 1, ..., dP, 1)
738738
interleaved_size = ones(Int, 2P)
739-
interleaved_size[1:2:2N] .= size(x)
739+
interleaved_size[1:2:(2N)] .= size(x)
740740

741741
x_interleaved = reshape(x, interleaved_size...)
742742

743743
# (d1, 1, d2, 1, ..., dP, 1) -> (d1, r1, d2, r2, ..., dP, rP)
744744
broadcast_target_size = interleaved_size
745-
broadcast_target_size[2:2:2M] .= counts
745+
broadcast_target_size[2:2:(2M)] .= counts
746746

747747
x_broadcasted = broadcast_to_size(x_interleaved, broadcast_target_size)
748748

749749
# (d1, r1, d2, r2, ..., dP, rP) -> (d1*r1, d2*r2, ..., dP*rP)
750-
final_size = vec(prod(reshape(broadcast_target_size, 2, :), dims=1))
750+
final_size = vec(prod(reshape(broadcast_target_size, 2, :); dims=1))
751751

752752
x_final = reshape(x_broadcasted, final_size...)
753753

test/basic.jl

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -366,8 +366,8 @@ end
366366

367367
@testset "repeat" begin
368368
@testset for (size, counts) in Iterators.product(
369-
[(2,), (2,3), (2,3,4), (2,3,4,5)],
370-
[(), (1,), (2,), (2,1), (1,2), (2,2), (2,2,2), (1,1,1,1,1)]
369+
[(2,), (2, 3), (2, 3, 4), (2, 3, 4, 5)],
370+
[(), (1,), (2,), (2, 1), (1, 2), (2, 2), (2, 2, 2), (1, 1, 1, 1, 1)],
371371
)
372372
x = rand(size...)
373373
@test (@jit repeat(Reactant.to_rarray(x), counts...)) == repeat(x, counts...)

0 commit comments

Comments
 (0)