Mercurial > repos > public > sbplib_julia

comparison src/LazyTensors/lazy_tensor_operations.jl @ 533:aac7cc1fa79a feature/inflated_tensormapping_transpose

Find changesets by keywords (author, files, the commit message), revision number or hash, or revset expression.
Try to improve the naming in split_index()
author Jonatan Werpers <jonatan@werpers.com>
date Thu, 26 Nov 2020 17:18:32 +0100
parents 588a843907de
children 41e82a5d4d48
comparison
equal deleted inserted replaced
532:588a843907de 533:aac7cc1fa79a
282 return apply_transpose(itm.tm, v_inner, inner_index...) 282 return apply_transpose(itm.tm, v_inner, inner_index...)
283 end 283 end
284 284
285 285
286 """ 286 """
287 split_index(::Val{A}, ::Val{B_view}, ::Val{B_middle}, ::Val{C}, I...) 287 split_index(::Val{dim_before}, ::Val{dim_view}, ::Val{dim_index}, ::Val{dim_after}, I...)
288 288
289 Splits the multi-index `I` into two parts. One part which is expected to be used as a view, and one which is expected to be used as an index. 289 Splits the multi-index `I` into two parts. One part which is expected to be
290 used as a view, and one which is expected to be used as an index.
290 Eg. 291 Eg.
291 ``` 292 ```
292 (1,2,3,4) -> (1,:,:,:,4), (2,3) 293 (1,2,3,4) -> (1,:,:,:,4), (2,3)
293 ``` 294 ```
294 295
295 `B_view` controls how many colons are in the view, and `B_middle` controls how many elements are extracted from the middle. 296 `dim_view` controls how many colons are in the view, and `dim_index` controls
296 `A` and `C` decides the length of the index parts before and after the colons in the view index. 297 how many elements are extracted from the middle.
297 298 `dim_before` and `dim_after` decides the length of the index parts before and after the colons in the view index.
298 Arguments should satisfy `length(I) == A+B_domain+C`. 299
300 Arguments should satisfy `length(I) == dim_before+B_domain+dim_after`.
299 301
300 The returned values satisfy 302 The returned values satisfy
301 * `length(view_index) == A + B_view + C` 303 * `length(view_index) == dim_before + dim_view + dim_after`
302 * `length(I_middle) == B_middle` 304 * `length(I_middle) == dim_index`
303 """ 305 """
304 function split_index(::Val{A}, ::Val{B_view}, ::Val{B_middle}, ::Val{C}, I...) where {A,B_view, B_middle,C} 306 function split_index(::Val{dim_before}, ::Val{dim_view}, ::Val{dim_index}, ::Val{dim_after}, I...) where {dim_before,dim_view, dim_index,dim_after}
305 I_before, I_middle, I_after = split_tuple(I, Val(A), Val(B_middle)) 307 I_before, I_middle, I_after = split_tuple(I, Val(dim_before), Val(dim_index))
306 308
307 view_index = (I_before..., ntuple((i)->:, B_view)..., I_after...) 309 view_index = (I_before..., ntuple((i)->:, dim_view)..., I_after...)
308 310
309 return view_index, I_middle 311 return view_index, I_middle
310 end 312 end
311 313
312 # TODO: Can this be replaced by something more elegant while still being type stable? 2020-10-21 314 # TODO: Can this be replaced by something more elegant while still being type stable? 2020-10-21
345 function split_tuple(t::NTuple{N},::Val{M},::Val{K}) where {N,M,K} 347 function split_tuple(t::NTuple{N},::Val{M},::Val{K}) where {N,M,K}
346 p1, tail = split_tuple(t, Val(M)) 348 p1, tail = split_tuple(t, Val(M))
347 p2, p3 = split_tuple(tail, Val(K)) 349 p2, p3 = split_tuple(tail, Val(K))
348 return p1,p2,p3 350 return p1,p2,p3
349 end 351 end
350
351 352
352 353
353 """ 354 """
354 flatten_tuple(t) 355 flatten_tuple(t)
355 356