Mercurial > repos > public > sbplib_julia

comparison src/LazyTensors/lazy_tensor_operations.jl @ 997:20c376dffe84 refactor/lazy_tensors

Find changesets by keywords (author, files, the commit message), revision number or hash, or revset expression.
Move tuple functions to their own file
author Jonatan Werpers <jonatan@werpers.com>
date Fri, 18 Mar 2022 21:26:02 +0100
parents 1ba8a398af9c
children 390dfc3db4b1
comparison
equal deleted inserted replaced
996:aa72f067e771 997:20c376dffe84
301 v_inner = view(v, view_index...) 301 v_inner = view(v, view_index...)
302 return apply_transpose(itm.tm, v_inner, inner_index...) 302 return apply_transpose(itm.tm, v_inner, inner_index...)
303 end 303 end
304 304
305 305
306 """
307 split_index(::Val{dim_before}, ::Val{dim_view}, ::Val{dim_index}, ::Val{dim_after}, I...)
308
309 Splits the multi-index `I` into two parts. One part which is expected to be
310 used as a view, and one which is expected to be used as an index.
311 Eg.
312 ```
313 split_index(Val(1),Val(3),Val(2),Val(1),(1,2,3,4)) -> (1,:,:,:,4), (2,3)
314 ```
315
316 `dim_view` controls how many colons are in the view, and `dim_index` controls
317 how many elements are extracted from the middle.
318 `dim_before` and `dim_after` decides the length of the index parts before and after the colons in the view index.
319
320 Arguments should satisfy `length(I) == dim_before+B_domain+dim_after`.
321
322 The returned values satisfy
323 * `length(view_index) == dim_before + dim_view + dim_after`
324 * `length(I_middle) == dim_index`
325 """
326 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}
327 I_before, I_middle, I_after = split_tuple(I, Val(dim_before), Val(dim_index))
328
329 view_index = (I_before..., ntuple((i)->:, dim_view)..., I_after...)
330
331 return view_index, I_middle
332 end
333
334 # TODO: Can this be replaced by something more elegant while still being type stable? 2020-10-21
335 # See:
336 # https://github.com/JuliaLang/julia/issues/34884
337 # https://github.com/JuliaLang/julia/issues/30386
338 """
339 slice_tuple(t, Val(l), Val(u))
340
341 Get a slice of a tuple in a type stable way.
342 Equivalent to `t[l:u]` but type stable.
343 """
344 function slice_tuple(t,::Val{L},::Val{U}) where {L,U}
345 return ntuple(i->t[i+L-1], U-L+1)
346 end
347
348 """
349 split_tuple(t::Tuple{...}, ::Val{M}) where {N,M}
350
351 Split the tuple `t` into two parts. the first part is `M` long.
352 E.g
353 ```julia
354 split_tuple((1,2,3,4),Val(3)) -> (1,2,3), (4,)
355 ```
356 """
357 function split_tuple(t::NTuple{N,Any},::Val{M}) where {N,M}
358 return slice_tuple(t,Val(1), Val(M)), slice_tuple(t,Val(M+1), Val(N))
359 end
360
361 """
362 split_tuple(t::Tuple{...},::Val{M},::Val{K}) where {N,M,K}
363
364 Same as `split_tuple(t::NTuple{N},::Val{M})` but splits the tuple in three parts. With the first
365 two parts having lenght `M` and `K`.
366 """
367 function split_tuple(t::NTuple{N,Any},::Val{M},::Val{K}) where {N,M,K}
368 p1, tail = split_tuple(t, Val(M))
369 p2, p3 = split_tuple(tail, Val(K))
370 return p1,p2,p3
371 end
372
373
374 """
375 flatten_tuple(t)
376
377 Takes a nested tuple and flattens the whole structure
378 """
379 flatten_tuple(t::NTuple{N, Number} where N) = t
380 flatten_tuple(t::Tuple) = ((flatten_tuple.(t)...)...,) # simplify?
381 flatten_tuple(ts::Vararg) = flatten_tuple(ts)
382
383 @doc raw""" 306 @doc raw"""
384 LazyOuterProduct(tms...) 307 LazyOuterProduct(tms...)
385 308
386 Creates a `LazyTensorComposition` for the outerproduct of `tms...`. 309 Creates a `LazyTensorComposition` for the outerproduct of `tms...`.
387 This is done by separating the outer product into regular products of outer products involving only identity mappings and one non-identity mapping. 310 This is done by separating the outer product into regular products of outer products involving only identity mappings and one non-identity mapping.