Merge pull request JuliaArrays#156 from ranocha/hr/fix_contiguous_bat…

…ch_size fix `contiguous_batch_size` for reshaped views
ranocha · May 31, 2021 · 6a8efa6 · 6a8efa6
2 parents 9789fa4 + ada0cf2
commit 6a8efa6
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 0 deletions.
diff --git a/src/stridelayout.jl b/src/stridelayout.jl
@@ -269,6 +269,7 @@ function _contiguous_batch_size(::StaticInt{D}, ::R) where {D,R<:Tuple}
         return nothing
     end
 end
+_contiguous_batch_size(::StaticInt{-1}, ::R) where {R<:Tuple} = -One()
 
 contiguous_batch_size(::Type{Array{T,N}}) where {T,N} = Zero()
 contiguous_batch_size(::Type{BitArray{N}}) where {N} = Zero()

diff --git a/test/runtests.jl b/test/runtests.jl
@@ -398,6 +398,12 @@ using OffsetArrays
     @test @inferred(contiguous_batch_size(@view(PermutedDimsArray(A,(3,1,2))[2:3,2,:]))) === ArrayInterface.StaticInt(-1)
     @test @inferred(contiguous_batch_size(@view(PermutedDimsArray(A,(3,1,2))[2:3,2,:])')) === ArrayInterface.StaticInt(-1)
     @test @inferred(contiguous_batch_size(@view(PermutedDimsArray(A,(3,1,2))[:,1:2,1])')) === ArrayInterface.StaticInt(0)
+    let u_base = randn(10, 10)
+        u_view = view(u_base, 3, :)
+        u_reshaped_view = reshape(u_view, 1, size(u_base, 2))
+        @test @inferred(contiguous_batch_size(u_view)) === ArrayInterface.StaticInt(-1)
+        @test @inferred(contiguous_batch_size(u_reshaped_view)) === ArrayInterface.StaticInt(-1)
+    end
 
     @test @inferred(stride_rank(@SArray(zeros(2,2,2)))) == (1, 2, 3)
     @test @inferred(stride_rank(A)) == (1,2,3)