changeset 1223:5bfb182e24dc refactor/LazyTensors/tuple_manipulation

Start adding simpler code
author Jonatan Werpers <jonatan@werpers.com>
date Sat, 18 Feb 2023 12:06:59 +0100
parents 7ee258e5289e
children e2f6dafb5d83
files src/LazyTensors/tuple_manipulation.jl test/LazyTensors/tuple_manipulation_test.jl
diffstat 2 files changed, 69 insertions(+), 0 deletions(-) [+]
line wrap: on
line diff
--- a/src/LazyTensors/tuple_manipulation.jl	Wed Feb 08 10:29:06 2023 +0100
+++ b/src/LazyTensors/tuple_manipulation.jl	Sat Feb 18 12:06:59 2023 +0100
@@ -25,6 +25,7 @@
 
     return view_index, I_middle
 end
+# TBD: If the nice split_tuple works, can this be cleaned up as well?
 
 # TODO: Can this be replaced by something more elegant while still being type stable? 2020-10-21
 # See:
@@ -65,6 +66,31 @@
     return p1,p2,p3
 end
 
+# TBD Are the above defs even needed? Can the below one be used without problems?
+
+"""
+    split_tuple(t, szs)
+
+Split the tuple `t` into a set of tuples of the sizes given in `szs`.
+`sum(szs)` should equal `lenght(t)`.
+"""
+function split_tuple(t, szs)
+    if length(t) != sum(szs; init=0)
+        throw(ArgumentError("length(t) must equal sum(szs)"))
+    end
+
+    rs = sizes_to_ranges(szs)
+    return map(r->t[r], rs)
+end
+
+function sizes_to_ranges(szs)
+    cum_szs = cumsum((0, szs...))
+    return ntuple(i->cum_szs[i]+1:cum_szs[i+1], length(szs))
+end
+
+
+concatenate_tuples(t::Tuple,ts::Vararg{Tuple}) = (t..., concatenate_tuples(ts...)...)
+concatenate_tuples(t::Tuple) = t
 
 """
     flatten_tuple(t)
@@ -74,6 +100,7 @@
 flatten_tuple(t::NTuple{N, Number} where N) = t
 flatten_tuple(t::Tuple) = ((flatten_tuple.(t)...)...,) # simplify?
 flatten_tuple(ts::Vararg) = flatten_tuple(ts)
+# TBD: Can concatenate_tuples be used instead?
 
 """
     left_pad_tuple(t, val, N)
--- a/test/LazyTensors/tuple_manipulation_test.jl	Wed Feb 08 10:29:06 2023 +0100
+++ b/test/LazyTensors/tuple_manipulation_test.jl	Sat Feb 18 12:06:59 2023 +0100
@@ -51,6 +51,48 @@
         @inferred LazyTensors.split_tuple((1,2,3,4,5,6),Val(3),Val(2))
         @inferred LazyTensors.split_tuple((1,true,3),Val(1), Val(1))
     end
+
+    @testset "general" begin
+        @test LazyTensors.split_tuple((),()) == ()
+        @test LazyTensors.split_tuple((),(0,)) == ((),)
+        @test LazyTensors.split_tuple((1,), (1,)) == tuple((1,))
+        @test LazyTensors.split_tuple((1,2), (1,1)) == tuple((1,),(2,))
+        @test LazyTensors.split_tuple((1,2), (0,1,1)) == tuple((),(1,),(2,))
+        @test LazyTensors.split_tuple((1,2), (1,0,1)) == tuple((1,),(),(2,))
+        @test LazyTensors.split_tuple((1,2), (1,1,0)) == tuple((1,),(2,),())
+        @test LazyTensors.split_tuple((1,2,3,4), (2,0,1,1)) == tuple((1,2),(),(3,),(4,))
+
+        err_msg = "length(t) must equal sum(szs)"
+        @test_throws ArgumentError(err_msg) LazyTensors.split_tuple((), (2,))
+        @test_throws ArgumentError(err_msg) LazyTensors.split_tuple((2,), ())
+        @test_throws ArgumentError(err_msg) LazyTensors.split_tuple((1,), (2,))
+        @test_throws ArgumentError(err_msg) LazyTensors.split_tuple((1,2), (1,2))
+        @test_throws ArgumentError(err_msg) LazyTensors.split_tuple((1,2), (1))
+
+
+        split_tuple_static(t, ::Val{SZS}) where {SZS} = @inline LazyTensors.split_tuple(t,SZS)
+
+        @inferred split_tuple_static((1,2,3,4,5,6), Val((3,1,2)))
+    end
+end
+
+@testset "sizes_to_ranges" begin
+    @test LazyTensors.sizes_to_ranges((1,)) == (1:1,)
+    @test LazyTensors.sizes_to_ranges((2,)) == (1:2,)
+    @test LazyTensors.sizes_to_ranges((2,3)) == (1:2,3:5)
+    @test LazyTensors.sizes_to_ranges((3,2,4)) == (1:3,4:5,6:9)
+    @test LazyTensors.sizes_to_ranges((0,2)) == (1:0,1:2)
+    @test LazyTensors.sizes_to_ranges((2,0)) == (1:2,2:1)
+    @test LazyTensors.sizes_to_ranges((2,0,3)) == (1:2,2:1,3:5)
+end
+
+@testset "concatenate_tuples" begin
+    @test LazyTensors.concatenate_tuples(()) == ()
+    @test LazyTensors.concatenate_tuples((1,)) == (1,)
+    @test LazyTensors.concatenate_tuples((1,), ()) == (1,)
+    @test LazyTensors.concatenate_tuples((),(1,)) == (1,)
+    @test LazyTensors.concatenate_tuples((1,2,3),(4,5)) == (1,2,3,4,5)
+    @test LazyTensors.concatenate_tuples((1,2,3),(4,5),(6,7)) == (1,2,3,4,5,6,7)
 end
 
 @testset "flatten_tuple" begin