Merge pull request #16 from CasBex/reduce_over_dims

_buffer_frame for Cartesian index
JuliaReinforcementLearning · Aug 9, 2023 · 1fb0f9c · 1fb0f9c
2 parents 45f79d4 + 9e08d85
commit 1fb0f9c
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 0 deletions.
diff --git a/src/CircularArrayBuffers.jl b/src/CircularArrayBuffers.jl
@@ -96,6 +96,7 @@ end
     end
 end
 
+_buffer_frame(cb::CircularArrayBuffer, I::CartesianIndex) = CartesianIndex(map(i->_buffer_frame(cb, i), Tuple(I)))
 _buffer_frame(cb::CircularArrayBuffer, I::AbstractArray{<:Integer}) = map(i -> _buffer_frame(cb, i), I)
 _buffer_frame(cb::CircularArrayBuffer, ::Colon) = _buffer_frame(cb, axes(cb)[end])
 

diff --git a/test/runtests.jl b/test/runtests.jl
@@ -129,6 +129,8 @@ CUDA.allowscalar(false)
             push!(b, x * A)
         end
 
+        @test sum(b, dims=3) == fill(6.0, 2, 2, 1)
+
         @test capacity(b) == 3
         @test isfull(b) == true
         @test length(b) == 2 * 2 * 3
@@ -176,6 +178,7 @@ CUDA.allowscalar(false)
         for i in 1:5
             push!(b, fill(i, 2))
         end
+        @test sum(b, dims=2) == [12;12;;]
         empty!(b)
         append!(b, 1:4)
         @test b == [
-Original file line number
+Diff line change
@@ Expand Up / @@ -96,6 +96,7 @@ end @@
         end
     end
+    _buffer_frame(cb::CircularArrayBuffer, I::CartesianIndex) = CartesianIndex(map(i->_buffer_frame(cb, i), Tuple(I)))
     _buffer_frame(cb::CircularArrayBuffer, I::AbstractArray{<:Integer}) = map(i -> _buffer_frame(cb, i), I)
     _buffer_frame(cb::CircularArrayBuffer, ::Colon) = _buffer_frame(cb, axes(cb)[end])
@@ Expand Down @@