add patch code for StackViews

Author: findmyway

src/patch.jl

@@ -1,3 +1,10 @@
 import MLUtils
 
-MLUtils.batch(x::AbstractArray{<:Number}) = x
+MLUtils.batch(x::AbstractArray{<:Number}) = x
+
+#####
+
+import StackViews: StackView
+
+lazy_stack(x) = StackView(x)
+lazy_stack(x::AbstractVector{<:Number}) = x

src/traces.jl

@@ -1,7 +1,6 @@
 export Trace, Traces, MultiplexTraces, Episode, Episodes
 
 import MacroTools: @forward
-import StackViews: StackView
 
 #####
 
@@ -196,7 +195,7 @@ end
 
 function Base.getindex(e::Episodes{names}, I) where {names}
     NamedTuple{names}(
-        StackView(
+        lazy_stack(
             map(I) do i
                 x, y = e.inds[i]
                 e.episodes[x][n][y]

test/samplers.jl

@@ -21,7 +21,7 @@
 
     @test length(e) == 5
     @test size(e[2:4].state) == (2, 3, 3)
-    @test_broken size(e[2:4].action) == (3,)
+    @test size(e[2:4].action) == (3,)
 end
 
 @testset "MetaSampler" begin

test/traces.jl

@@ -167,12 +167,12 @@ end
 
     @test t[end] == (state=2.0, action=2)
 
-    # https://github.com/JuliaArrays/StackViews.jl/issues/3
-    @test_broken t[1:2] == (state=[1.0, 2.0], action=[1, 2])
+    @test t[1:2] == (state=[1.0, 2.0], action=[1, 2])
 
     push!(t, (state=3.0, action=3))
     t[] = true # seal
 
+    # a vector of episode-level partitions is returned for now
     @test_broken size(t[:state]) == (3,)
 
     push!(t, Episode(Traces(state=[4.0, 5.0, 6.0], action=[4, 5, 6])))