Mercurial > repos > public > sbplib_julia

comparison src/LazyTensors/lazy_tensor_operations.jl @ 1044:f857057e61e6 refactor/sbpoperators/inflation

Find changesets by keywords (author, files, the commit message), revision number or hash, or revset expression.
Merge default
author Jonatan Werpers <jonatan@werpers.com>
date Tue, 22 Mar 2022 22:05:34 +0100
parents 52f07c77299d 9e76bf19904c
children 5a3281429a48 423a6442efc3 62f321caa964
comparison
equal deleted inserted replaced
1024:5be17f647018 1044:f857057e61e6
1 """ 1 """
2 LazyTensorApplication{T,R,D} <: LazyArray{T,R} 2 TensorApplication{T,R,D} <: LazyArray{T,R}
3 3
4 Struct for lazy application of a LazyTensor. Created using `*`. 4 Struct for lazy application of a LazyTensor. Created using `*`.
5 5
6 Allows the result of a `LazyTensor` applied to a vector to be treated as an `AbstractArray`. 6 Allows the result of a `LazyTensor` applied to a vector to be treated as an `AbstractArray`.
7 With a mapping `m` and a vector `v` the LazyTensorApplication object can be created by `m*v`. 7 With a mapping `m` and a vector `v` the TensorApplication object can be created by `m*v`.
8 The actual result will be calcualted when indexing into `m*v`. 8 The actual result will be calcualted when indexing into `m*v`.
9 """ 9 """
10 struct LazyTensorApplication{T,R,D, TM<:LazyTensor{<:Any,R,D}, AA<:AbstractArray{<:Any,D}} <: LazyArray{T,R} 10 struct TensorApplication{T,R,D, TM<:LazyTensor{<:Any,R,D}, AA<:AbstractArray{<:Any,D}} <: LazyArray{T,R}
11 t::TM 11 t::TM
12 o::AA 12 o::AA
13 13
14 function LazyTensorApplication(t::LazyTensor{<:Any,R,D}, o::AbstractArray{<:Any,D}) where {R,D} 14 function TensorApplication(t::LazyTensor{<:Any,R,D}, o::AbstractArray{<:Any,D}) where {R,D}
15 @boundscheck check_domain_size(t, size(o)) 15 @boundscheck check_domain_size(t, size(o))
16 I = ntuple(i->1, range_dim(t)) 16 I = ntuple(i->1, range_dim(t))
17 T = typeof(apply(t,o,I...)) 17 T = typeof(apply(t,o,I...))
18 return new{T,R,D,typeof(t), typeof(o)}(t,o) 18 return new{T,R,D,typeof(t), typeof(o)}(t,o)
19 end 19 end
20 end 20 end
21 21
22 function Base.getindex(ta::LazyTensorApplication{T,R}, I::Vararg{Any,R}) where {T,R} 22 function Base.getindex(ta::TensorApplication{T,R}, I::Vararg{Any,R}) where {T,R}
23 @boundscheck checkbounds(ta, Int.(I)...) 23 @boundscheck checkbounds(ta, Int.(I)...)
24 return apply(ta.t, ta.o, I...) 24 return @inbounds apply(ta.t, ta.o, I...)
25 end 25 end
26 Base.getindex(ta::LazyTensorApplication{T,1} where T, I::CartesianIndex{1}) = ta[Tuple(I)...] # Would otherwise be caught in the previous method. 26 Base.@propagate_inbounds Base.getindex(ta::TensorApplication{T,1} where T, I::CartesianIndex{1}) = ta[Tuple(I)...] # Would otherwise be caught in the previous method.
27 Base.size(ta::LazyTensorApplication) = range_size(ta.t) 27 Base.size(ta::TensorApplication) = range_size(ta.t)
28 28
29 29
30 """ 30 """
31 LazyTensorTranspose{T,R,D} <: LazyTensor{T,D,R} 31 TensorTranspose{T,R,D} <: LazyTensor{T,D,R}
32 32
33 Struct for lazy transpose of a LazyTensor. 33 Struct for lazy transpose of a LazyTensor.
34 34
35 If a mapping implements the the `apply_transpose` method this allows working with 35 If a mapping implements the the `apply_transpose` method this allows working with
36 the transpose of mapping `m` by using `m'`. `m'` will work as a regular LazyTensor lazily calling 36 the transpose of mapping `m` by using `m'`. `m'` will work as a regular LazyTensor lazily calling
37 the appropriate methods of `m`. 37 the appropriate methods of `m`.
38 """ 38 """
39 struct LazyTensorTranspose{T,R,D, TM<:LazyTensor{T,R,D}} <: LazyTensor{T,D,R} 39 struct TensorTranspose{T,R,D, TM<:LazyTensor{T,R,D}} <: LazyTensor{T,D,R}
40 tm::TM 40 tm::TM
41 end 41 end
42 42
43 # # TBD: Should this be implemented on a type by type basis or through a trait to provide earlier errors? 43 # # TBD: Should this be implemented on a type by type basis or through a trait to provide earlier errors?
44 # Jonatan 2020-09-25: Is the problem that you can take the transpose of any LazyTensor even if it doesn't implement `apply_transpose`? 44 # Jonatan 2020-09-25: Is the problem that you can take the transpose of any LazyTensor even if it doesn't implement `apply_transpose`?
45 Base.adjoint(tm::LazyTensor) = LazyTensorTranspose(tm) 45 Base.adjoint(tm::LazyTensor) = TensorTranspose(tm)
46 Base.adjoint(tmt::LazyTensorTranspose) = tmt.tm 46 Base.adjoint(tmt::TensorTranspose) = tmt.tm
47 47
48 apply(tmt::LazyTensorTranspose{T,R,D}, v::AbstractArray{<:Any,R}, I::Vararg{Any,D}) where {T,R,D} = apply_transpose(tmt.tm, v, I...) 48 apply(tmt::TensorTranspose{T,R,D}, v::AbstractArray{<:Any,R}, I::Vararg{Any,D}) where {T,R,D} = apply_transpose(tmt.tm, v, I...)
49 apply_transpose(tmt::LazyTensorTranspose{T,R,D}, v::AbstractArray{<:Any,D}, I::Vararg{Any,R}) where {T,R,D} = apply(tmt.tm, v, I...) 49 apply_transpose(tmt::TensorTranspose{T,R,D}, v::AbstractArray{<:Any,D}, I::Vararg{Any,R}) where {T,R,D} = apply(tmt.tm, v, I...)
50 50
51 range_size(tmt::LazyTensorTranspose) = domain_size(tmt.tm) 51 range_size(tmt::TensorTranspose) = domain_size(tmt.tm)
52 domain_size(tmt::LazyTensorTranspose) = range_size(tmt.tm) 52 domain_size(tmt::TensorTranspose) = range_size(tmt.tm)
53 53
54 54
55 struct LazyTensorBinaryOperation{Op,T,R,D,T1<:LazyTensor{T,R,D},T2<:LazyTensor{T,R,D}} <: LazyTensor{T,D,R} 55 struct ElementwiseTensorOperation{Op,T,R,D,T1<:LazyTensor{T,R,D},T2<:LazyTensor{T,R,D}} <: LazyTensor{T,D,R}
56 tm1::T1 56 tm1::T1
57 tm2::T2 57 tm2::T2
58 58
59 function LazyTensorBinaryOperation{Op,T,R,D}(tm1::T1,tm2::T2) where {Op,T,R,D, T1<:LazyTensor{T,R,D},T2<:LazyTensor{T,R,D}} 59 function ElementwiseTensorOperation{Op,T,R,D}(tm1::T1,tm2::T2) where {Op,T,R,D, T1<:LazyTensor{T,R,D},T2<:LazyTensor{T,R,D}}
60 @boundscheck check_domain_size(tm2, domain_size(tm1)) 60 @boundscheck check_domain_size(tm2, domain_size(tm1))
61 @boundscheck check_range_size(tm2, range_size(tm1)) 61 @boundscheck check_range_size(tm2, range_size(tm1))
62 return new{Op,T,R,D,T1,T2}(tm1,tm2) 62 return new{Op,T,R,D,T1,T2}(tm1,tm2)
63 end 63 end
64 end 64 end
65 65
66 LazyTensorBinaryOperation{Op}(s,t) where Op = LazyTensorBinaryOperation{Op,eltype(s), range_dim(s), domain_dim(s)}(s,t) 66 ElementwiseTensorOperation{Op}(s,t) where Op = ElementwiseTensorOperation{Op,eltype(s), range_dim(s), domain_dim(s)}(s,t)
67 67
68 apply(tmBinOp::LazyTensorBinaryOperation{:+,T,R,D}, v::AbstractArray{<:Any,D}, I::Vararg{Any,R}) where {T,R,D} = apply(tmBinOp.tm1, v, I...) + apply(tmBinOp.tm2, v, I...) 68 apply(tmBinOp::ElementwiseTensorOperation{:+,T,R,D}, v::AbstractArray{<:Any,D}, I::Vararg{Any,R}) where {T,R,D} = apply(tmBinOp.tm1, v, I...) + apply(tmBinOp.tm2, v, I...)
69 apply(tmBinOp::LazyTensorBinaryOperation{:-,T,R,D}, v::AbstractArray{<:Any,D}, I::Vararg{Any,R}) where {T,R,D} = apply(tmBinOp.tm1, v, I...) - apply(tmBinOp.tm2, v, I...) 69 apply(tmBinOp::ElementwiseTensorOperation{:-,T,R,D}, v::AbstractArray{<:Any,D}, I::Vararg{Any,R}) where {T,R,D} = apply(tmBinOp.tm1, v, I...) - apply(tmBinOp.tm2, v, I...)
70 70
71 range_size(tmBinOp::LazyTensorBinaryOperation) = range_size(tmBinOp.tm1) 71 range_size(tmBinOp::ElementwiseTensorOperation) = range_size(tmBinOp.tm1)
72 domain_size(tmBinOp::LazyTensorBinaryOperation) = domain_size(tmBinOp.tm1) 72 domain_size(tmBinOp::ElementwiseTensorOperation) = domain_size(tmBinOp.tm1)
73 73
74 74
75 """ 75 """
76 LazyTensorComposition{T,R,K,D} 76 TensorComposition{T,R,K,D}
77 77
78 Lazily compose two `LazyTensor`s, so that they can be handled as a single `LazyTensor`. 78 Lazily compose two `LazyTensor`s, so that they can be handled as a single `LazyTensor`.
79 """ 79 """
80 struct LazyTensorComposition{T,R,K,D, TM1<:LazyTensor{T,R,K}, TM2<:LazyTensor{T,K,D}} <: LazyTensor{T,R,D} 80 struct TensorComposition{T,R,K,D, TM1<:LazyTensor{T,R,K}, TM2<:LazyTensor{T,K,D}} <: LazyTensor{T,R,D}
81 t1::TM1 81 t1::TM1
82 t2::TM2 82 t2::TM2
83 83
84 function LazyTensorComposition(t1::LazyTensor{T,R,K}, t2::LazyTensor{T,K,D}) where {T,R,K,D} 84 function TensorComposition(t1::LazyTensor{T,R,K}, t2::LazyTensor{T,K,D}) where {T,R,K,D}
85 @boundscheck check_domain_size(t1, range_size(t2)) 85 @boundscheck check_domain_size(t1, range_size(t2))
86 return new{T,R,K,D, typeof(t1), typeof(t2)}(t1,t2) 86 return new{T,R,K,D, typeof(t1), typeof(t2)}(t1,t2)
87 end 87 end
88 end 88 end
89 89
90 range_size(tm::LazyTensorComposition) = range_size(tm.t1) 90 range_size(tm::TensorComposition) = range_size(tm.t1)
91 domain_size(tm::LazyTensorComposition) = domain_size(tm.t2) 91 domain_size(tm::TensorComposition) = domain_size(tm.t2)
92 92
93 function apply(c::LazyTensorComposition{T,R,K,D}, v::AbstractArray{<:Any,D}, I::Vararg{Any,R}) where {T,R,K,D} 93 function apply(c::TensorComposition{T,R,K,D}, v::AbstractArray{<:Any,D}, I::Vararg{Any,R}) where {T,R,K,D}
94 apply(c.t1, c.t2*v, I...) 94 apply(c.t1, c.t2*v, I...)
95 end 95 end
96 96
97 function apply_transpose(c::LazyTensorComposition{T,R,K,D}, v::AbstractArray{<:Any,R}, I::Vararg{Any,D}) where {T,R,K,D} 97 function apply_transpose(c::TensorComposition{T,R,K,D}, v::AbstractArray{<:Any,R}, I::Vararg{Any,D}) where {T,R,K,D}
98 apply_transpose(c.t2, c.t1'*v, I...) 98 apply_transpose(c.t2, c.t1'*v, I...)
99 end 99 end
100 100
101 101
102 """ 102 """
103 LazyTensorComposition(tm, tmi::IdentityTensor) 103 TensorComposition(tm, tmi::IdentityTensor)
104 LazyTensorComposition(tmi::IdentityTensor, tm) 104 TensorComposition(tmi::IdentityTensor, tm)
105 105
106 Composes a `Tensormapping` `tm` with an `IdentityTensor` `tmi`, by returning `tm` 106 Composes a `Tensormapping` `tm` with an `IdentityTensor` `tmi`, by returning `tm`
107 """ 107 """
108 function LazyTensorComposition(tm::LazyTensor{T,R,D}, tmi::IdentityTensor{T,D}) where {T,R,D} 108 function TensorComposition(tm::LazyTensor{T,R,D}, tmi::IdentityTensor{T,D}) where {T,R,D}
109 @boundscheck check_domain_size(tm, range_size(tmi)) 109 @boundscheck check_domain_size(tm, range_size(tmi))
110 return tm 110 return tm
111 end 111 end
112 112
113 function LazyTensorComposition(tmi::IdentityTensor{T,R}, tm::LazyTensor{T,R,D}) where {T,R,D} 113 function TensorComposition(tmi::IdentityTensor{T,R}, tm::LazyTensor{T,R,D}) where {T,R,D}
114 @boundscheck check_domain_size(tmi, range_size(tm)) 114 @boundscheck check_domain_size(tmi, range_size(tm))
115 return tm 115 return tm
116 end 116 end
117 # Specialization for the case where tm is an IdentityTensor. Required to resolve ambiguity. 117 # Specialization for the case where tm is an IdentityTensor. Required to resolve ambiguity.
118 function LazyTensorComposition(tm::IdentityTensor{T,D}, tmi::IdentityTensor{T,D}) where {T,D} 118 function TensorComposition(tm::IdentityTensor{T,D}, tmi::IdentityTensor{T,D}) where {T,D}
119 @boundscheck check_domain_size(tm, range_size(tmi)) 119 @boundscheck check_domain_size(tm, range_size(tmi))
120 return tmi 120 return tmi
121 end 121 end
122 122
123 123
124 """ 124 """
125 InflatedLazyTensor{T,R,D} <: LazyTensor{T,R,D} 125 InflatedTensor{T,R,D} <: LazyTensor{T,R,D}
126 126
127 An inflated `LazyTensor` with dimensions added before and afer its actual dimensions. 127 An inflated `LazyTensor` with dimensions added before and afer its actual dimensions.
128 """ 128 """
129 struct InflatedLazyTensor{T,R,D,D_before,R_middle,D_middle,D_after, TM<:LazyTensor{T,R_middle,D_middle}} <: LazyTensor{T,R,D} 129 struct InflatedTensor{T,R,D,D_before,R_middle,D_middle,D_after, TM<:LazyTensor{T,R_middle,D_middle}} <: LazyTensor{T,R,D}
130 before::IdentityTensor{T,D_before} 130 before::IdentityTensor{T,D_before}
131 tm::TM 131 tm::TM
132 after::IdentityTensor{T,D_after} 132 after::IdentityTensor{T,D_after}
133 133
134 function InflatedLazyTensor(before, tm::LazyTensor{T}, after) where T 134 function InflatedTensor(before, tm::LazyTensor{T}, after) where T
135 R_before = range_dim(before) 135 R_before = range_dim(before)
136 R_middle = range_dim(tm) 136 R_middle = range_dim(tm)
137 R_after = range_dim(after) 137 R_after = range_dim(after)
138 R = R_before+R_middle+R_after 138 R = R_before+R_middle+R_after
139 139
144 return new{T,R,D,D_before,R_middle,D_middle,D_after, typeof(tm)}(before, tm, after) 144 return new{T,R,D,D_before,R_middle,D_middle,D_after, typeof(tm)}(before, tm, after)
145 end 145 end
146 end 146 end
147 147
148 """ 148 """
149 InflatedLazyTensor(before, tm, after) 149 InflatedTensor(before, tm, after)
150 InflatedLazyTensor(before,tm) 150 InflatedTensor(before,tm)
151 InflatedLazyTensor(tm,after) 151 InflatedTensor(tm,after)
152 152
153 The outer product of `before`, `tm` and `after`, where `before` and `after` are `IdentityTensor`s. 153 The outer product of `before`, `tm` and `after`, where `before` and `after` are `IdentityTensor`s.
154 154
155 If one of `before` or `after` is left out, a 0-dimensional `IdentityTensor` is used as the default value. 155 If one of `before` or `after` is left out, a 0-dimensional `IdentityTensor` is used as the default value.
156 156
157 If `tm` already is an `InflatedLazyTensor`, `before` and `after` will be extended instead of 157 If `tm` already is an `InflatedTensor`, `before` and `after` will be extended instead of
158 creating a nested `InflatedLazyTensor`. 158 creating a nested `InflatedTensor`.
159 """ 159 """
160 InflatedLazyTensor(::IdentityTensor, ::LazyTensor, ::IdentityTensor) 160 InflatedTensor(::IdentityTensor, ::LazyTensor, ::IdentityTensor)
161 161
162 function InflatedLazyTensor(before, itm::InflatedLazyTensor, after) 162 function InflatedTensor(before, itm::InflatedTensor, after)
163 return InflatedLazyTensor( 163 return InflatedTensor(
164 IdentityTensor(before.size..., itm.before.size...), 164 IdentityTensor(before.size..., itm.before.size...),
165 itm.tm, 165 itm.tm,
166 IdentityTensor(itm.after.size..., after.size...), 166 IdentityTensor(itm.after.size..., after.size...),
167 ) 167 )
168 end 168 end
169 169
170 InflatedLazyTensor(before::IdentityTensor, tm::LazyTensor{T}) where T = InflatedLazyTensor(before,tm,IdentityTensor{T}()) 170 InflatedTensor(before::IdentityTensor, tm::LazyTensor{T}) where T = InflatedTensor(before,tm,IdentityTensor{T}())
171 InflatedLazyTensor(tm::LazyTensor{T}, after::IdentityTensor) where T = InflatedLazyTensor(IdentityTensor{T}(),tm,after) 171 InflatedTensor(tm::LazyTensor{T}, after::IdentityTensor) where T = InflatedTensor(IdentityTensor{T}(),tm,after)
172 # Resolve ambiguity between the two previous methods 172 # Resolve ambiguity between the two previous methods
173 InflatedLazyTensor(I1::IdentityTensor{T}, I2::IdentityTensor{T}) where T = InflatedLazyTensor(I1,I2,IdentityTensor{T}()) 173 InflatedTensor(I1::IdentityTensor{T}, I2::IdentityTensor{T}) where T = InflatedTensor(I1,I2,IdentityTensor{T}())
174 174
175 # TODO: Implement some pretty printing in terms of ⊗. E.g InflatedLazyTensor(I(3),B,I(2)) -> I(3)⊗B⊗I(2) 175 # TODO: Implement some pretty printing in terms of ⊗. E.g InflatedTensor(I(3),B,I(2)) -> I(3)⊗B⊗I(2)
176 176
177 function range_size(itm::InflatedLazyTensor) 177 function range_size(itm::InflatedTensor)
178 return flatten_tuple( 178 return flatten_tuple(
179 range_size(itm.before), 179 range_size(itm.before),
180 range_size(itm.tm), 180 range_size(itm.tm),
181 range_size(itm.after), 181 range_size(itm.after),
182 ) 182 )
183 end 183 end
184 184
185 function domain_size(itm::InflatedLazyTensor) 185 function domain_size(itm::InflatedTensor)
186 return flatten_tuple( 186 return flatten_tuple(
187 domain_size(itm.before), 187 domain_size(itm.before),
188 domain_size(itm.tm), 188 domain_size(itm.tm),
189 domain_size(itm.after), 189 domain_size(itm.after),
190 ) 190 )
191 end 191 end
192 192
193 function apply(itm::InflatedLazyTensor{T,R,D}, v::AbstractArray{<:Any,D}, I::Vararg{Any,R}) where {T,R,D} 193 function apply(itm::InflatedTensor{T,R,D}, v::AbstractArray{<:Any,D}, I::Vararg{Any,R}) where {T,R,D}
194 dim_before = range_dim(itm.before) 194 dim_before = range_dim(itm.before)
195 dim_domain = domain_dim(itm.tm) 195 dim_domain = domain_dim(itm.tm)
196 dim_range = range_dim(itm.tm) 196 dim_range = range_dim(itm.tm)
197 dim_after = range_dim(itm.after) 197 dim_after = range_dim(itm.after)
198 198
200 200
201 v_inner = view(v, view_index...) 201 v_inner = view(v, view_index...)
202 return apply(itm.tm, v_inner, inner_index...) 202 return apply(itm.tm, v_inner, inner_index...)
203 end 203 end
204 204
205 function apply_transpose(itm::InflatedLazyTensor{T,R,D}, v::AbstractArray{<:Any,R}, I::Vararg{Any,D}) where {T,R,D} 205 function apply_transpose(itm::InflatedTensor{T,R,D}, v::AbstractArray{<:Any,R}, I::Vararg{Any,D}) where {T,R,D}
206 dim_before = range_dim(itm.before) 206 dim_before = range_dim(itm.before)
207 dim_domain = domain_dim(itm.tm) 207 dim_domain = domain_dim(itm.tm)
208 dim_range = range_dim(itm.tm) 208 dim_range = range_dim(itm.tm)
209 dim_after = range_dim(itm.after) 209 dim_after = range_dim(itm.after)
210 210
216 216
217 217
218 @doc raw""" 218 @doc raw"""
219 LazyOuterProduct(tms...) 219 LazyOuterProduct(tms...)
220 220
221 Creates a `LazyTensorComposition` for the outerproduct of `tms...`. 221 Creates a `TensorComposition` for the outerproduct of `tms...`.
222 This is done by separating the outer product into regular products of outer products involving only identity mappings and one non-identity mapping. 222 This is done by separating the outer product into regular products of outer products involving only identity mappings and one non-identity mapping.
223 223
224 First let 224 First let
225 ```math 225 ```math
226 \begin{aligned} 226 \begin{aligned}
253 ``` 253 ```
254 """ 254 """
255 function LazyOuterProduct end 255 function LazyOuterProduct end
256 256
257 function LazyOuterProduct(tm1::LazyTensor{T}, tm2::LazyTensor{T}) where T 257 function LazyOuterProduct(tm1::LazyTensor{T}, tm2::LazyTensor{T}) where T
258 itm1 = InflatedLazyTensor(tm1, IdentityTensor{T}(range_size(tm2))) 258 itm1 = InflatedTensor(tm1, IdentityTensor{T}(range_size(tm2)))
259 itm2 = InflatedLazyTensor(IdentityTensor{T}(domain_size(tm1)),tm2) 259 itm2 = InflatedTensor(IdentityTensor{T}(domain_size(tm1)),tm2)
260 260
261 return itm1∘itm2 261 return itm1∘itm2
262 end 262 end
263 263
264 LazyOuterProduct(t1::IdentityTensor{T}, t2::IdentityTensor{T}) where T = IdentityTensor{T}(t1.size...,t2.size...) 264 LazyOuterProduct(t1::IdentityTensor{T}, t2::IdentityTensor{T}) where T = IdentityTensor{T}(t1.size...,t2.size...)
265 LazyOuterProduct(t1::LazyTensor, t2::IdentityTensor) = InflatedLazyTensor(t1, t2) 265 LazyOuterProduct(t1::LazyTensor, t2::IdentityTensor) = InflatedTensor(t1, t2)
266 LazyOuterProduct(t1::IdentityTensor, t2::LazyTensor) = InflatedLazyTensor(t1, t2) 266 LazyOuterProduct(t1::IdentityTensor, t2::LazyTensor) = InflatedTensor(t1, t2)
267 267
268 LazyOuterProduct(tms::Vararg{LazyTensor}) = foldl(LazyOuterProduct, tms) 268 LazyOuterProduct(tms::Vararg{LazyTensor}) = foldl(LazyOuterProduct, tms)
269 269
270 270
271 271