Skip to content

Commit 137bef2

Browse files
authored
Preliminary support for qr with padded matrices and use layout_broadcasted for all LazyArrays (#379)
* Preliminary support for qr with padded matrices * generalise layout_broadcasted * Add more layout_broadcasted * Fix padded tests * Update cache.jl * fix cachedarray bug * increase coverage * simplify ZerosLayout overloads * increase cov
1 parent 6ff9502 commit 137bef2

File tree

11 files changed

+251
-105
lines changed

11 files changed

+251
-105
lines changed

Project.toml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
name = "LazyArrays"
22
uuid = "5078a376-72f3-5289-bfd5-ec5146d43c02"
3-
version = "2.6.4"
3+
version = "2.7"
44

55
[deps]
66
ArrayLayouts = "4c555306-a7a7-4459-81d9-ec55ddd5c99a"

src/LazyArrays.jl

Lines changed: 17 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -39,7 +39,8 @@ import ArrayLayouts: AbstractQLayout, Dot, Dotu, Ldiv, Lmul, MatMulMatAdd, MatMu
3939
materialize!, mulreduce, reshapedlayout, rowsupport, scalarone, scalarzero, sub_materialize,
4040
sublayout, symmetriclayout, symtridiagonallayout, transposelayout, triangulardata,
4141
triangularlayout, tridiagonallayout, zero!, transtype, OnesLayout,
42-
diagonaldata, subdiagonaldata, supdiagonaldata, MemoryLayout
42+
diagonaldata, subdiagonaldata, supdiagonaldata, MemoryLayout, MatLmulVec, MatLmulMat,
43+
AdjQRCompactWYQLayout
4344

4445
import FillArrays: AbstractFill, getindex_value
4546

@@ -49,6 +50,21 @@ export Mul, Applied, MulArray, MulVector, MulMatrix, InvMatrix, PInvMatrix,
4950
PaddedArray, PaddedVector, PaddedMatrix
5051

5152

53+
54+
"""
55+
LazyArray(x::Applied) :: ApplyArray
56+
LazyArray(x::Broadcasted) :: BroadcastArray
57+
58+
Wrap a lazy object that wraps a computation producing an array to an
59+
array.
60+
"""
61+
abstract type LazyArray{T,N} <: LayoutArray{T,N} end
62+
63+
const LazyMatrix{T} = LazyArray{T,2}
64+
const LazyVector{T} = LazyArray{T,1}
65+
66+
67+
5268
include("lazyapplying.jl")
5369
include("lazybroadcasting.jl")
5470
include("linalg/linalg.jl")

src/cache.jl

Lines changed: 11 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -336,10 +336,10 @@ MemoryLayout(C::Type{CachedArray{T,N,DAT,ARR}}) where {T,N,DAT,ARR} = cachedlayo
336336
BroadcastStyle(::Type{<:CachedArray{<:Any,N}}) where N = LazyArrayStyle{N}()
337337

338338
broadcasted(::LazyArrayStyle, op, A::CachedArray) = CachedArray(broadcast(op, cacheddata(A)), broadcast(op, A.array))
339-
broadcasted(::LazyArrayStyle, op, A::CachedArray, c::Number) = CachedArray(broadcast(op, cacheddata(A), c), broadcast(op, A.array, c))
340-
broadcasted(::LazyArrayStyle, op, c::Number, A::CachedArray) = CachedArray(broadcast(op, c, cacheddata(A)), broadcast(op, c, A.array))
341-
broadcasted(::LazyArrayStyle, op, A::CachedArray, c::Ref) = CachedArray(broadcast(op, cacheddata(A), c), broadcast(op, A.array, c))
342-
broadcasted(::LazyArrayStyle, op, c::Ref, A::CachedArray) = CachedArray(broadcast(op, c, cacheddata(A)), broadcast(op, c, A.array))
339+
layout_broadcasted(::CachedLayout, _, op, A::AbstractArray, c::Number) = CachedArray(broadcast(op, cacheddata(A), c), broadcast(op, A.array, c))
340+
layout_broadcasted(_, ::CachedLayout, op, c::Number, A::CachedArray) = CachedArray(broadcast(op, c, cacheddata(A)), broadcast(op, c, A.array))
341+
layout_broadcasted(::CachedLayout, _, op, A::CachedArray, c::Ref) = CachedArray(broadcast(op, cacheddata(A), c), broadcast(op, A.array, c))
342+
layout_broadcasted(_, ::CachedLayout, op, c::Ref, A::CachedArray) = CachedArray(broadcast(op, c, cacheddata(A)), broadcast(op, c, A.array))
343343

344344

345345
function layout_broadcasted(::CachedLayout, _, op, A::AbstractVector, B::AbstractVector)
@@ -365,28 +365,16 @@ function layout_broadcasted(::CachedLayout, ::CachedLayout, op, A::AbstractVecto
365365
CachedArray(convert(Array, broadcast(op, Adat, Bdat)), broadcast(op, A.array, B.array))
366366
end
367367

368-
function layout_broadcasted(op, A, B)
369-
if axes(A) axes(B)
370-
(size(A,1) == 1 || size(B,1) == 1) && error("Internal error: Scalar-like broadcasting not yet supported.")
371-
throw(DimensionMismatch("arrays could not be broadcast to a common size; got a dimension with lengths $(length(A)) and $(length(B))"))
372-
end
373-
layout_broadcasted(MemoryLayout(A), MemoryLayout(B), op, A, B)
374-
end
375368

376-
broadcasted(::LazyArrayStyle, op, A::CachedVector, B::AbstractVector) = layout_broadcasted(op, A, B)
377-
broadcasted(::LazyArrayStyle, op, A::AbstractVector, B::CachedVector) = layout_broadcasted(op, A, B)
378-
broadcasted(::LazyArrayStyle, op, A::CachedVector, B::CachedVector) = layout_broadcasted(op, A, B)
379-
broadcasted(::LazyArrayStyle, op, A::Broadcasted, B::CachedVector) = broadcast(op, materialize(A), B)
380369

381-
broadcasted(::LazyArrayStyle, op, A::SubArray{<:Any,1,<:CachedMatrix}, B::CachedVector) = layout_broadcasted(op, A, B)
382-
broadcasted(::LazyArrayStyle, op, A::SubArray{<:Any,1,<:CachedMatrix}, B::AbstractVector) = layout_broadcasted(op, A, B)
383-
broadcasted(::LazyArrayStyle, op, A::CachedVector, B::SubArray{<:Any,1,<:CachedMatrix}) = layout_broadcasted(op, A, B)
384-
broadcasted(::LazyArrayStyle, op, A::AbstractVector, B::SubArray{<:Any,1,<:CachedMatrix}) = layout_broadcasted(op, A, B)
385370

386-
broadcasted(::LazyArrayStyle{1}, op, a::CachedVector, b::Zeros{<:Any,1}) = broadcast(DefaultArrayStyle{1}(), op, a, b)
387-
broadcasted(::LazyArrayStyle{1}, op, a::Zeros{<:Any,1}, b::CachedVector) = broadcast(DefaultArrayStyle{1}(), op, a, b)
388-
broadcasted(::LazyArrayStyle{1}, ::typeof(*), a::CachedVector, b::Zeros{<:Any,1}) = broadcast(DefaultArrayStyle{1}(), *, a, b)
389-
broadcasted(::LazyArrayStyle{1}, ::typeof(*), a::Zeros{<:Any,1}, b::CachedVector) = broadcast(DefaultArrayStyle{1}(), *, a, b)
371+
for op in (:*, :/, :+, :-)
372+
@eval layout_broadcasted(::CachedLayout, ::ZerosLayout, ::typeof($op), a::AbstractVector, b::AbstractVector) = broadcast(DefaultArrayStyle{1}(), $op, a, b)
373+
end
374+
375+
for op in (:*, :\, :+, :-)
376+
@eval layout_broadcasted(::ZerosLayout, ::CachedLayout, ::typeof($op), a::AbstractVector, b::AbstractVector) = broadcast(DefaultArrayStyle{1}(), $op, a, b)
377+
end
390378

391379

392380

src/lazyapplying.jl

Lines changed: 0 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -180,17 +180,6 @@ eltype(A::Applied{<:MatrixFunctionStyle}) = float(eltype(first(A.args)))
180180

181181
getindex(A::Applied, kj...) = materialize(A)[kj...]
182182

183-
"""
184-
LazyArray(x::Applied) :: ApplyArray
185-
LazyArray(x::Broadcasted) :: BroadcastArray
186-
187-
Wrap a lazy object that wraps a computation producing an array to an
188-
array.
189-
"""
190-
abstract type LazyArray{T,N} <: LayoutArray{T,N} end
191-
192-
const LazyMatrix{T} = LazyArray{T,2}
193-
const LazyVector{T} = LazyArray{T,1}
194183

195184
struct ApplyArray{T, N, F, Args<:Tuple} <: LazyArray{T,N}
196185
f::F

src/lazybroadcasting.jl

Lines changed: 17 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,23 @@
11
struct LazyArrayStyle{N} <: AbstractArrayStyle{N} end
22
LazyArrayStyle(::Val{N}) where N = LazyArrayStyle{N}()
33
LazyArrayStyle{M}(::Val{N}) where {N,M} = LazyArrayStyle{N}()
4+
5+
6+
7+
layout_broadcasted(_, _, op, A, B) = Base.Broadcast.Broadcasted(Base.Broadcast.combine_styles(A,B), op, (A, B))
8+
layout_broadcasted(op, A, B) = layout_broadcasted(MemoryLayout(A), MemoryLayout(B), op, A, B)
9+
10+
DefaultArrayStyle(::LazyArrayStyle{N}) where N = DefaultArrayStyle{N}()
11+
broadcasted(::LazyArrayStyle, op, A, B) = layout_broadcasted(op, A, B)
12+
13+
for op in (:*, :/, :+, :-)
14+
@eval layout_broadcasted(::ZerosLayout, _, ::typeof($op), a, b) = broadcasted(DefaultArrayStyle(Base.Broadcast.combine_styles(a,b)), $op, a, b)
15+
end
16+
17+
for op in (:*, :\, :+, :-)
18+
@eval layout_broadcasted(_, ::ZerosLayout, ::typeof($op), a, b) = broadcasted(DefaultArrayStyle(Base.Broadcast.combine_styles(a,b)), $op, a, b)
19+
end
20+
421
"""
522
BroadcastLayout{F}()
623
@@ -174,22 +191,6 @@ broadcasted(::LazyArrayStyle{1}, ::typeof(*), a::AbstractRange, b::AbstractFill)
174191
broadcasted(::LazyArrayStyle{1}, ::typeof(*), a::Zeros{<:Any,1}, b::AbstractRange) = broadcast(DefaultArrayStyle{1}(), *, a, b)
175192
broadcasted(::LazyArrayStyle{1}, ::typeof(*), a::AbstractRange, b::Zeros{<:Any,1}) = broadcast(DefaultArrayStyle{1}(), *, a, b)
176193

177-
for op in (:*, :/, :\)
178-
@eval broadcasted(::LazyArrayStyle{N}, ::typeof($op), a::Zeros{T,N}, b::Zeros{V,N}) where {T,V,N} = broadcast(DefaultArrayStyle{N}(), $op, a, b)
179-
end
180-
181-
for op in (:*, :/)
182-
@eval begin
183-
broadcasted(::LazyArrayStyle{N}, ::typeof($op), a::Zeros{T,N}, b::AbstractArray{V,N}) where {T,V,N} = broadcast(DefaultArrayStyle{N}(), $op, a, b)
184-
broadcasted(::LazyArrayStyle{N}, ::typeof($op), a::Zeros{T,N}, b::Broadcasted) where {T,N} = broadcast(DefaultArrayStyle{N}(), $op, a, b)
185-
end
186-
end
187-
for op in (:*, :\)
188-
@eval begin
189-
broadcasted(::LazyArrayStyle{N}, ::typeof($op), a::AbstractArray{T,N}, b::Zeros{V,N}) where {T,V,N} = broadcast(DefaultArrayStyle{N}(), $op, a, b)
190-
broadcasted(::LazyArrayStyle{N}, ::typeof($op), a::Broadcasted, b::Zeros{V,N}) where {V,N} = broadcast(DefaultArrayStyle{N}(), $op, a, b)
191-
end
192-
end
193194

194195
###
195196
# support

src/lazyconcat.jl

Lines changed: 34 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -455,12 +455,24 @@ broadcasted(::LazyArrayStyle, op, A::Transpose{<:Any,<:Vcat}) = transpose(broadc
455455
broadcasted(::LazyArrayStyle, op, A::Adjoint{<:Real,<:Vcat}) = broadcast(op, parent(A))'
456456

457457

458-
for Cat in (:Vcat, :Hcat)
458+
for Cat in (:vcat, :hcat)
459459
@eval begin
460-
broadcasted(::LazyArrayStyle, op, A::$Cat, c::Number) = $Cat(_flatten_nums(A.args, broadcast((x,y) -> broadcast(op, x, y), A.args, c))...)
461-
broadcasted(::LazyArrayStyle, op, c::Number, A::$Cat) = $Cat(_flatten_nums(A.args, broadcast((x,y) -> broadcast(op, x, y), c, A.args))...)
462-
broadcasted(::LazyArrayStyle, op, A::$Cat, c::Ref) = $Cat(_flatten_nums(A.args, broadcast((x,y) -> broadcast(op, x, Ref(y)), A.args, c))...)
463-
broadcasted(::LazyArrayStyle, op, c::Ref, A::$Cat) = $Cat(_flatten_nums(A.args, broadcast((x,y) -> broadcast(op, Ref(x), y), c, A.args))...)
460+
function layout_broadcasted(Alay::ApplyLayout{typeof($Cat)}, _, op, A::AbstractArray, c::Number)
461+
Aargs = arguments(Alay, A)
462+
ApplyArray($Cat, _flatten_nums(Aargs, broadcast((x,y) -> broadcast(op, x, y), Aargs, c))...)
463+
end
464+
function layout_broadcasted(_, Alay::ApplyLayout{typeof($Cat)}, op, c::Number, A::AbstractArray)
465+
Aargs = arguments(Alay, A)
466+
ApplyArray($Cat, _flatten_nums(Aargs, broadcast((x,y) -> broadcast(op, x, y), c, Aargs))...)
467+
end
468+
function layout_broadcasted(Alay::ApplyLayout{typeof($Cat)}, _, op, A::AbstractArray, c::Ref)
469+
Aargs = arguments(Alay, A)
470+
ApplyArray($Cat, _flatten_nums(Aargs, broadcast((x,y) -> broadcast(op, x, Ref(y)), Aargs, c))...)
471+
end
472+
function layout_broadcasted(_, Alay::ApplyLayout{typeof($Cat)}, op, c::Ref, A::AbstractArray)
473+
Aargs = arguments(Alay, A)
474+
ApplyArray($Cat, _flatten_nums(Aargs, broadcast((x,y) -> broadcast(op, Ref(x), y), c, Aargs))...)
475+
end
464476
end
465477
end
466478

@@ -483,16 +495,27 @@ layout_broadcasted(::AbstractLazyLayout, ::ApplyLayout{typeof(vcat)}, op, A::Abs
483495
layout_broadcasted(::ApplyLayout{typeof(vcat)}, lay::CachedLayout, op, A::AbstractVector, B::AbstractVector) = layout_broadcasted(UnknownLayout(), lay, op, A, B)
484496
layout_broadcasted(lay::CachedLayout, ::ApplyLayout{typeof(vcat)}, op, A::AbstractVector, B::AbstractVector) = layout_broadcasted(lay, UnknownLayout(), op, A, B)
485497

486-
function layout_broadcasted(::ApplyLayout{typeof(vcat)}, _, op, A::AbstractVector, B::AbstractVector)
487-
kr = _vcat_axes(map(axes,A.args)...) # determine how to break up B
498+
for op in (:*, :/, :+, :-)
499+
@eval layout_broadcasted(::ZerosLayout, ::ApplyLayout{typeof(vcat)}, ::typeof($op), a::AbstractVector, b::AbstractVector) = layout_broadcasted(ZerosLayout(), UnknownLayout(), $op, a, b)
500+
end
501+
for op in (:*, :\, :+, :-)
502+
@eval layout_broadcasted(::ApplyLayout{typeof(vcat)}, ::ZerosLayout, ::typeof($op), a::AbstractVector, b::AbstractVector) = layout_broadcasted(UnknownLayout(), ZerosLayout(), $op, a, b)
503+
end
504+
505+
506+
507+
function layout_broadcasted(Alay::ApplyLayout{typeof(vcat)}, _, op, A::AbstractVector, B::AbstractVector)
508+
Aargs = arguments(Alay, A)
509+
kr = _vcat_axes(map(axes, Aargs)...) # determine how to break up B
488510
B_arrays = _vcat_getindex_eval(B,kr...) # evaluate B at same chunks as A
489-
ApplyVector(vcat, broadcast((a,b) -> broadcast(op,a,b), A.args, B_arrays)...)
511+
ApplyVector(vcat, broadcast((a,b) -> broadcast(op,a,b), Aargs, B_arrays)...)
490512
end
491513

492-
function layout_broadcasted(_, ::ApplyLayout{typeof(vcat)}, op, A::AbstractVector, B::AbstractVector)
493-
kr = _vcat_axes(axes.(B.args)...)
514+
function layout_broadcasted(_, Blay::ApplyLayout{typeof(vcat)}, op, A::AbstractVector, B::AbstractVector)
515+
Bargs = arguments(Blay, B)
516+
kr = _vcat_axes(axes.(Bargs)...)
494517
A_arrays = _vcat_getindex_eval(A,kr...)
495-
Vcat(broadcast((a,b) -> broadcast(op,a,b), A_arrays, B.args)...)
518+
Vcat(broadcast((a,b) -> broadcast(op,a,b), A_arrays, Bargs)...)
496519
end
497520

498521

@@ -554,24 +577,6 @@ end
554577
_vcat_layout_broadcasted((Ahead,Atail)::Tuple{Number,Any}, (Bhead,Btail)::Tuple{Number,Any}, op, A, B) = Vcat(op.(Ahead,Bhead), op.(Atail,Btail))
555578

556579

557-
broadcasted(::LazyArrayStyle, op, a::Vcat{<:Any,N}, b::AbstractArray{<:Any,N}) where N = layout_broadcasted(op, a, b)
558-
broadcasted(::LazyArrayStyle, op, a::AbstractArray{<:Any,N}, b::Vcat{<:Any,N}) where N = layout_broadcasted(op, a, b)
559-
broadcasted(::LazyArrayStyle{1}, op, a::Vcat{<:Any,1}, b::Zeros{<:Any,1}) = broadcast(DefaultArrayStyle{1}(), op, a, b)
560-
broadcasted(::LazyArrayStyle{1}, op, a::Zeros{<:Any,1}, b::Vcat{<:Any,1}) = broadcast(DefaultArrayStyle{1}(), op, a, b)
561-
broadcasted(::LazyArrayStyle{1}, ::typeof(\), a::Vcat{<:Any,1}, b::Zeros{<:Any,1}) = broadcast(DefaultArrayStyle{1}(), \, a, b)
562-
broadcasted(::LazyArrayStyle{1}, ::typeof(/), a::Zeros{<:Any,1}, b::Vcat{<:Any,1}) = broadcast(DefaultArrayStyle{1}(), /, a, b)
563-
broadcasted(::LazyArrayStyle{1}, ::typeof(*), a::Vcat{<:Any,1}, b::Zeros{<:Any,1}) = broadcast(DefaultArrayStyle{1}(), *, a, b)
564-
broadcasted(::LazyArrayStyle{1}, ::typeof(*), a::Zeros{<:Any,1}, b::Vcat{<:Any,1}) = broadcast(DefaultArrayStyle{1}(), *, a, b)
565-
566-
567-
# Cannot broadcast Vcat's in a lazy way so stick to BroadcastArray
568-
broadcasted(::LazyArrayStyle, op, A::Vcat, B::Vcat) = layout_broadcasted(op, A, B)
569-
570-
# ambiguities
571-
broadcasted(::LazyArrayStyle, op, A::Vcat{<:Any,1}, B::CachedVector) = layout_broadcasted(op, A, B)
572-
broadcasted(::LazyArrayStyle, op, A::CachedVector, B::Vcat{<:Any,1}) = layout_broadcasted(op, A, B)
573-
574-
575580

576581
function +(A::Vcat, B::Vcat)
577582
size(A) == size(B) || throw(DimensionMismatch("dimensions must match."))

0 commit comments

Comments
 (0)