Mercurial > repos > public > sbplib_julia

comparison src/LazyTensors/lazy_tensor_operations.jl @ 455:b86312d14873 feature/inflated_tensormapping

Find changesets by keywords (author, files, the commit message), revision number or hash, or revset expression.
Make split_index type stable
author Jonatan Werpers <jonatan@werpers.com>
date Wed, 21 Oct 2020 16:29:59 +0200
parents aeda2698166d
children 8f4c31e06689
comparison
equal deleted inserted replaced
454:eb4c34438e30 455:b86312d14873
229 (1,2,3,4) -> (1,:,:,4), (2,3) 229 (1,2,3,4) -> (1,:,:,4), (2,3)
230 ``` 230 ```
231 """ 231 """
232 function split_index(itm::InflatedTensorMapping{T,R,D}, I::Vararg{Any,R}) where {T,R,D} 232 function split_index(itm::InflatedTensorMapping{T,R,D}, I::Vararg{Any,R}) where {T,R,D}
233 I_before = I[1:range_dim(itm.before)] 233 I_before = I[1:range_dim(itm.before)]
234 I_after = I[(end-range_dim(itm.after)+1):end] 234 I_after = slice_tuple(I,Val(R-range_dim(itm.after)+1),Val(R))
235 235
236 view_index = (I_before..., ntuple((i)->:,domain_dim(itm.tm))..., I_after...) 236 view_index = (I_before..., ntuple((i)->:,domain_dim(itm.tm))..., I_after...)
237 inner_index = I[range_dim(itm.before)+1:end-range_dim(itm.after)] 237 inner_index = slice_tuple(I, Val(range_dim(itm.before)+1), Val(R-range_dim(itm.after)))
238 238
239 return (view_index, inner_index) 239 return (view_index, inner_index)
240 end
241
242 function slice_tuple(t,::Val{L},::Val{U}) where {L,U}
243 return ntuple(i->t[i+L-1], U-L+1)
240 end 244 end
241 245
242 flatten_tuple(t::NTuple{N, Number} where N) = t 246 flatten_tuple(t::NTuple{N, Number} where N) = t
243 flatten_tuple(t::Tuple) = ((flatten_tuple.(t)...)...,) # simplify? 247 flatten_tuple(t::Tuple) = ((flatten_tuple.(t)...)...,) # simplify?
244 flatten_tuple(ts::Vararg) = flatten_tuple(ts) 248 flatten_tuple(ts::Vararg) = flatten_tuple(ts)