Implement

Author: DanielVandH

Project.toml

@@ -1,6 +1,6 @@
 name = "LazyArrays"
 uuid = "5078a376-72f3-5289-bfd5-ec5146d43c02"
-version = "2.9.2"
+version = "2.9.3"
 
 [deps]
 ArrayLayouts = "4c555306-a7a7-4459-81d9-ec55ddd5c99a"

src/LazyArrays.jl

@@ -18,7 +18,7 @@ import Base: *, +, -, /, <, ==, >, \, ≤, ≥, (:), @_gc_preserve_begin, @_gc_p
              oneto, add_sum, promote_op
 
 import Base.Broadcast: AbstractArrayStyle, BroadcastStyle, Broadcasted, DefaultArrayStyle, broadcasted, combine_eltypes,
-                       instantiate
+                       instantiate, result_style, Unknown
 
 import LinearAlgebra: AbstractQ, AdjOrTrans, StructuredMatrixStyle, checksquare, det, diag, dot, lmul!, logabsdet,
                       norm1, norm2, normInf, normp, pinv, rmul!, tr, tril, triu

src/lazyconcat.jl

@@ -439,8 +439,16 @@ end
 # to take advantage of special implementations of the sub-components
 ######
 
-BroadcastStyle(::Type{<:Vcat{<:Any,N}}) where N = LazyArrayStyle{N}()
-BroadcastStyle(::Type{<:Hcat{<:Any}}) = LazyArrayStyle{2}()
+@inline tuple_type_broadcastlayout(::Type{I}) where {I<:Tuple} = result_style(BroadcastStyle(Base.tuple_type_head(I)), tuple_type_broadcastlayout(Base.tuple_type_tail(I)))
+@inline tuple_type_broadcastlayout(::Type{Tuple{}}) = Unknown()
+@inline tuple_type_broadcastlayout(::Type{Tuple{A}}) where {A} = BroadcastStyle(A)
+@inline tuple_type_broadcastlayout(::Type{Tuple{A,B}}) where {A,B} = result_style(BroadcastStyle(A), BroadcastStyle(B))
+@inline tuple_type_broadcastlayout(::Type{Tuple{A,B,C}}) where {A,B,C} = result_style(BroadcastStyle(A), tuple_type_broadcastlayout(Tuple{B,C}))
+@inline tuple_type_broadcastlayout(::Type{Tuple{A,B,C,D}}) where {A,B,C,D} = result_style(BroadcastStyle(A), tuple_type_broadcastlayout(Tuple{B,C,D}))
+@inline tuple_type_broadcastlayout(::Type{Tuple{A,B,C,D,E}}) where {A,B,C,D,E} = result_style(BroadcastStyle(A), tuple_type_broadcastlayout(Tuple{B,C,D,E}))
+
+BroadcastStyle(::Type{<:Vcat{<:Any,N,I}}) where {N,I<:Tuple} = result_style(LazyArrayStyle{N}(), tuple_type_broadcastlayout(I)) # the <:Tuple is to avoid ambiguity 
+BroadcastStyle(::Type{<:Hcat{<:Any,I}}) where {I<:Tuple} = result_style(LazyArrayStyle{2}(), tuple_type_broadcastlayout(I))
 
 # This is if we broadcast a function on a mixed concat f.([1; [2,3]])
 # such that f returns a vector, e.g., f(1) == [1,2], we don't want

test/cachetests.jl

@@ -6,6 +6,7 @@ import LazyArrays: CachedArray, CachedMatrix, CachedVector, PaddedLayout, Cached
                     CachedAbstractArray, CachedAbstractVector, CachedAbstractMatrix, AbstractCachedArray, AbstractCachedMatrix,
                     PaddedColumns, cacheddata, LazyArrayStyle, maybe_cacheddata, Accumulate, CachedArrayStyle, GenericCachedLayout,
                     AccumulateAbstractVector
+import Base.Broadcast: BroadcastStyle
 
 using ..InfiniteArrays
 using .InfiniteArrays: OneToInf
@@ -574,6 +575,18 @@ using Infinities
         @test cacheddata(F) === view(cacheddata(parent(parent(F))), 1:1)'
         @test cacheddata(G) === adjoint(view(cacheddata(parent(G)), 1:1, 1:1))
     end
+
+    @testset "BroadcastStyle for Vcat/Hcat of CachedArrayStyles" begin
+        @test BroadcastStyle(typeof(Vcat(cache(1:3), cache(4:6)))) == CachedArrayStyle{1}()
+        d = Accumulate(*, 1:10)
+        @test BroadcastStyle(typeof(Vcat(d, d))) == CachedArrayStyle{1}()
+        @test BroadcastStyle(typeof(Vcat(d', d'))) == CachedArrayStyle{2}()
+        @test BroadcastStyle(typeof(Hcat(d, d))) == CachedArrayStyle{2}()
+        @test BroadcastStyle(typeof(Vcat(d', (1:10)'))) == CachedArrayStyle{2}()
+        @test BroadcastStyle(typeof(Vcat((1:10)', d'))) == CachedArrayStyle{2}()
+        @test BroadcastStyle(typeof(Hcat(d, (1:10)))) == CachedArrayStyle{2}()
+        @test BroadcastStyle(typeof(Hcat((1:10), d))) == CachedArrayStyle{2}()
+    end
 end
 
 end # module

test/concattests.jl

@@ -4,7 +4,9 @@ using LazyArrays, FillArrays, LinearAlgebra, ArrayLayouts, Test, Base64
 using StaticArrays
 import LazyArrays: MemoryLayout, DenseColumnMajor, materialize!, call, paddeddata,
                     MulAdd, Applied, ApplyLayout, DefaultApplyStyle, sub_materialize, resizedata!,
-                    CachedVector, ApplyLayout, arguments, BroadcastVector, LazyLayout, cacheddata
+                    CachedVector, ApplyLayout, arguments, BroadcastVector, LazyLayout, cacheddata,
+                    LazyArrayStyle, CachedArrayStyle, Accumulate
+import Base.Broadcast: BroadcastStyle
 
 @testset "concat" begin
     @testset "Vcat" begin
@@ -712,4 +714,21 @@ import LazyArrays: MemoryLayout, DenseColumnMajor, materialize!, call, paddeddat
     end
 end
 
+@testset "BroadcastStyle" begin
+    args = (1:10, Accumulate(*, 1:10), BroadcastVector(exp, 1:10), BroadcastMatrix(exp, rand(10, 2)), Vcat(Accumulate(*, 1:10)', (1:10)')', [1, 2, 3, 4, 5, 6, 7, 8, 9, 10])
+    for i in 1:6 
+        @test @inferred(LazyArrays.tuple_type_broadcastlayout(typeof(args[1:i]))) == Base.Broadcast.combine_styles(args[1:i]...)
+    end
+    @test @inferred(LazyArrays.tuple_type_broadcastlayout(typeof(args))) == Base.Broadcast.combine_styles(args...) == CachedArrayStyle{2}()
+    @test @inferred(LazyArrays.tuple_type_broadcastlayout(Tuple{})) == Base.Broadcast.Unknown()
+    @test BroadcastStyle(typeof(Vcat(adjoint.(args)...))) == BroadcastStyle(typeof(Vcat(transpose.(args)...))) == CachedArrayStyle{2}()
+    @test BroadcastStyle(typeof(Hcat(args...))) == CachedArrayStyle{2}()
+    
+    @test BroadcastStyle(typeof(Vcat((1:10)'))) == LazyArrayStyle{2}() # make sure we preserve Lazy even without lazy args  
+    @test BroadcastStyle(typeof(Hcat((1:10)'))) == LazyArrayStyle{2}()
+    
+    @test BroadcastStyle(typeof(Vcat())) == LazyArrayStyle{1}() 
+    @test BroadcastStyle(typeof(Hcat())) == LazyArrayStyle{2}()
+end
+
 end # module