comparison src/LazyTensors/lazy_tensor_operations.jl @ 1100:157a78959e5d refactor/sbpoperators/inflation

Bring up to date
author Jonatan Werpers <jonatan@werpers.com>
date Tue, 10 May 2022 20:34:20 +0200
parents 2278730f9cee
children 6f51160c7ca7 f1c2a4fa0ee1
comparison
equal deleted inserted replaced
1099:05a25a5063bb 1100:157a78959e5d
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 TensorComposition(tm, tmi::IdentityTensor)
103 LazyTensorComposition(tm, tmi::IdentityTensor) 103 TensorComposition(tmi::IdentityTensor, tm)
104 LazyTensorComposition(tmi::IdentityTensor, tm)
105 104
106 Composes a `Tensormapping` `tm` with an `IdentityTensor` `tmi`, by returning `tm` 105 Composes a `Tensormapping` `tm` with an `IdentityTensor` `tmi`, by returning `tm`
107 """ 106 """
108 function LazyTensorComposition(tm::LazyTensor{T,R,D}, tmi::IdentityTensor{T,D}) where {T,R,D} 107 function TensorComposition(tm::LazyTensor{T,R,D}, tmi::IdentityTensor{T,D}) where {T,R,D}
109 @boundscheck check_domain_size(tm, range_size(tmi)) 108 @boundscheck check_domain_size(tm, range_size(tmi))
110 return tm 109 return tm
111 end 110 end
112 111
113 function LazyTensorComposition(tmi::IdentityTensor{T,R}, tm::LazyTensor{T,R,D}) where {T,R,D} 112 function TensorComposition(tmi::IdentityTensor{T,R}, tm::LazyTensor{T,R,D}) where {T,R,D}
114 @boundscheck check_domain_size(tmi, range_size(tm)) 113 @boundscheck check_domain_size(tmi, range_size(tm))
115 return tm 114 return tm
116 end 115 end
117 # Specialization for the case where tm is an IdentityTensor. Required to resolve ambiguity. 116 # 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} 117 function TensorComposition(tm::IdentityTensor{T,D}, tmi::IdentityTensor{T,D}) where {T,D}
119 @boundscheck check_domain_size(tm, range_size(tmi)) 118 @boundscheck check_domain_size(tm, range_size(tmi))
120 return tmi 119 return tmi
121 end 120 end
122 121
123 122 Base.:*(a::T, tm::LazyTensor{T}) where T = TensorComposition(ScalingTensor{T,range_dim(tm)}(a,range_size(tm)), tm)
124 """ 123 Base.:*(tm::LazyTensor{T}, a::T) where T = a*tm
125 InflatedLazyTensor{T,R,D} <: LazyTensor{T,R,D} 124
125 """
126 InflatedTensor{T,R,D} <: LazyTensor{T,R,D}
126 127
127 An inflated `LazyTensor` with dimensions added before and afer its actual dimensions. 128 An inflated `LazyTensor` with dimensions added before and afer its actual dimensions.
128 """ 129 """
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} 130 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} 131 before::IdentityTensor{T,D_before}
131 tm::TM 132 tm::TM
132 after::IdentityTensor{T,D_after} 133 after::IdentityTensor{T,D_after}
133 134
134 function InflatedLazyTensor(before, tm::LazyTensor{T}, after) where T 135 function InflatedTensor(before, tm::LazyTensor{T}, after) where T
135 R_before = range_dim(before) 136 R_before = range_dim(before)
136 R_middle = range_dim(tm) 137 R_middle = range_dim(tm)
137 R_after = range_dim(after) 138 R_after = range_dim(after)
138 R = R_before+R_middle+R_after 139 R = R_before+R_middle+R_after
139 140
144 return new{T,R,D,D_before,R_middle,D_middle,D_after, typeof(tm)}(before, tm, after) 145 return new{T,R,D,D_before,R_middle,D_middle,D_after, typeof(tm)}(before, tm, after)
145 end 146 end
146 end 147 end
147 148
148 """ 149 """
149 InflatedLazyTensor(before, tm, after) 150 InflatedTensor(before, tm, after)
150 InflatedLazyTensor(before,tm) 151 InflatedTensor(before,tm)
151 InflatedLazyTensor(tm,after) 152 InflatedTensor(tm,after)
152 153
153 The outer product of `before`, `tm` and `after`, where `before` and `after` are `IdentityTensor`s. 154 The outer product of `before`, `tm` and `after`, where `before` and `after` are `IdentityTensor`s.
154 155
155 If one of `before` or `after` is left out, a 0-dimensional `IdentityTensor` is used as the default value. 156 If one of `before` or `after` is left out, a 0-dimensional `IdentityTensor` is used as the default value.
156 157
157 If `tm` already is an `InflatedLazyTensor`, `before` and `after` will be extended instead of 158 If `tm` already is an `InflatedTensor`, `before` and `after` will be extended instead of
158 creating a nested `InflatedLazyTensor`. 159 creating a nested `InflatedTensor`.
159 """ 160 """
160 InflatedLazyTensor(::IdentityTensor, ::LazyTensor, ::IdentityTensor) 161 InflatedTensor(::IdentityTensor, ::LazyTensor, ::IdentityTensor)
161 162
162 function InflatedLazyTensor(before, itm::InflatedLazyTensor, after) 163 function InflatedTensor(before, itm::InflatedTensor, after)
163 return InflatedLazyTensor( 164 return InflatedTensor(
164 IdentityTensor(before.size..., itm.before.size...), 165 IdentityTensor(before.size..., itm.before.size...),
165 itm.tm, 166 itm.tm,
166 IdentityTensor(itm.after.size..., after.size...), 167 IdentityTensor(itm.after.size..., after.size...),
167 ) 168 )
168 end 169 end
169 170
170 InflatedLazyTensor(before::IdentityTensor, tm::LazyTensor{T}) where T = InflatedLazyTensor(before,tm,IdentityTensor{T}()) 171 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) 172 InflatedTensor(tm::LazyTensor{T}, after::IdentityTensor) where T = InflatedTensor(IdentityTensor{T}(),tm,after)
172 # Resolve ambiguity between the two previous methods 173 # Resolve ambiguity between the two previous methods
173 InflatedLazyTensor(I1::IdentityTensor{T}, I2::IdentityTensor{T}) where T = InflatedLazyTensor(I1,I2,IdentityTensor{T}()) 174 InflatedTensor(I1::IdentityTensor{T}, I2::IdentityTensor{T}) where T = InflatedTensor(I1,I2,IdentityTensor{T}())
174 175
175 # TODO: Implement some pretty printing in terms of ⊗. E.g InflatedLazyTensor(I(3),B,I(2)) -> I(3)⊗B⊗I(2) 176 # TODO: Implement some pretty printing in terms of ⊗. E.g InflatedTensor(I(3),B,I(2)) -> I(3)⊗B⊗I(2)
176 177
177 function range_size(itm::InflatedLazyTensor) 178 function range_size(itm::InflatedTensor)
178 return flatten_tuple( 179 return flatten_tuple(
179 range_size(itm.before), 180 range_size(itm.before),
180 range_size(itm.tm), 181 range_size(itm.tm),
181 range_size(itm.after), 182 range_size(itm.after),
182 ) 183 )
183 end 184 end
184 185
185 function domain_size(itm::InflatedLazyTensor) 186 function domain_size(itm::InflatedTensor)
186 return flatten_tuple( 187 return flatten_tuple(
187 domain_size(itm.before), 188 domain_size(itm.before),
188 domain_size(itm.tm), 189 domain_size(itm.tm),
189 domain_size(itm.after), 190 domain_size(itm.after),
190 ) 191 )
191 end 192 end
192 193
193 function apply(itm::InflatedLazyTensor{T,R,D}, v::AbstractArray{<:Any,D}, I::Vararg{Any,R}) where {T,R,D} 194 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) 195 dim_before = range_dim(itm.before)
195 dim_domain = domain_dim(itm.tm) 196 dim_domain = domain_dim(itm.tm)
196 dim_range = range_dim(itm.tm) 197 dim_range = range_dim(itm.tm)
197 dim_after = range_dim(itm.after) 198 dim_after = range_dim(itm.after)
198 199
200 201
201 v_inner = view(v, view_index...) 202 v_inner = view(v, view_index...)
202 return apply(itm.tm, v_inner, inner_index...) 203 return apply(itm.tm, v_inner, inner_index...)
203 end 204 end
204 205
205 function apply_transpose(itm::InflatedLazyTensor{T,R,D}, v::AbstractArray{<:Any,R}, I::Vararg{Any,D}) where {T,R,D} 206 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) 207 dim_before = range_dim(itm.before)
207 dim_domain = domain_dim(itm.tm) 208 dim_domain = domain_dim(itm.tm)
208 dim_range = range_dim(itm.tm) 209 dim_range = range_dim(itm.tm)
209 dim_after = range_dim(itm.after) 210 dim_after = range_dim(itm.after)
210 211
216 217
217 218
218 @doc raw""" 219 @doc raw"""
219 LazyOuterProduct(tms...) 220 LazyOuterProduct(tms...)
220 221
221 Creates a `LazyTensorComposition` for the outerproduct of `tms...`. 222 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. 223 This is done by separating the outer product into regular products of outer products involving only identity mappings and one non-identity mapping.
223 224
224 First let 225 First let
225 ```math 226 ```math
226 \begin{aligned} 227 \begin{aligned}
253 ``` 254 ```
254 """ 255 """
255 function LazyOuterProduct end 256 function LazyOuterProduct end
256 257
257 function LazyOuterProduct(tm1::LazyTensor{T}, tm2::LazyTensor{T}) where T 258 function LazyOuterProduct(tm1::LazyTensor{T}, tm2::LazyTensor{T}) where T
258 itm1 = InflatedLazyTensor(tm1, IdentityTensor{T}(range_size(tm2))) 259 itm1 = InflatedTensor(tm1, IdentityTensor{T}(range_size(tm2)))
259 itm2 = InflatedLazyTensor(IdentityTensor{T}(domain_size(tm1)),tm2) 260 itm2 = InflatedTensor(IdentityTensor{T}(domain_size(tm1)),tm2)
260 261
261 return itm1∘itm2 262 return itm1∘itm2
262 end 263 end
263 264
264 LazyOuterProduct(t1::IdentityTensor{T}, t2::IdentityTensor{T}) where T = IdentityTensor{T}(t1.size...,t2.size...) 265 LazyOuterProduct(t1::IdentityTensor{T}, t2::IdentityTensor{T}) where T = IdentityTensor{T}(t1.size...,t2.size...)
265 LazyOuterProduct(t1::LazyTensor, t2::IdentityTensor) = InflatedLazyTensor(t1, t2) 266 LazyOuterProduct(t1::LazyTensor, t2::IdentityTensor) = InflatedTensor(t1, t2)
266 LazyOuterProduct(t1::IdentityTensor, t2::LazyTensor) = InflatedLazyTensor(t1, t2) 267 LazyOuterProduct(t1::IdentityTensor, t2::LazyTensor) = InflatedTensor(t1, t2)
267 268
268 LazyOuterProduct(tms::Vararg{LazyTensor}) = foldl(LazyOuterProduct, tms) 269 LazyOuterProduct(tms::Vararg{LazyTensor}) = foldl(LazyOuterProduct, tms)
269 270
270 271
271 272
272 """ 273 """
273 inflate(tm, sz, dir) 274 inflate(tm::LazyTensor, sz, dir)
274 275
275 Inflate `tm` with identity tensors in all directions `d` for `d != dir`. 276 Inflate `tm` such that it gets the size `sz` in all directions except `dir`.
276 277 Here `sz[dir]` is ignored and replaced with the range and domains size of
277 # TODO: Describe when it is useful 278 `tm`.
279
280 An example of when this operation is useful is when extending a one
281 dimensional difference operator `D` to a 2D grid of a ceratin size. In that
282 case we could have
283
284 ```julia
285 Dx = inflate(D, (10,10), 1)
286 Dy = inflate(D, (10,10), 2)
287 ```
278 """ 288 """
279 function inflate(tm::LazyTensor, sz, dir) 289 function inflate(tm::LazyTensor, sz, dir)
280 Is = IdentityTensor{eltype(tm)}.(sz) 290 Is = IdentityTensor{eltype(tm)}.(sz)
281 parts = Base.setindex(Is, tm, dir) 291 parts = Base.setindex(Is, tm, dir)
282 return foldl(⊗, parts) 292 return foldl(⊗, parts)