From 5f6a64660b16d8f188caca5216e55debf4264611 Mon Sep 17 00:00:00 2001 From: Tom Smeding Date: Mon, 10 Jun 2024 23:28:11 +0200 Subject: Add *flatten --- src/Data/Array/Nested/Internal/Shaped.hs | 5 +++++ 1 file changed, 5 insertions(+) (limited to 'src/Data/Array/Nested/Internal/Shaped.hs') diff --git a/src/Data/Array/Nested/Internal/Shaped.hs b/src/Data/Array/Nested/Internal/Shaped.hs index 863e604..1855015 100644 --- a/src/Data/Array/Nested/Internal/Shaped.hs +++ b/src/Data/Array/Nested/Internal/Shaped.hs @@ -370,6 +370,11 @@ srev1 arr = slift (sshape arr) (\_ -> X.rev1) arr sreshape :: Elt a => ShS sh' -> Shaped sh a -> Shaped sh' a sreshape sh' (Shaped arr) = Shaped (mreshape (shCvtSX sh') arr) +sflatten :: Elt a => Shaped sh a -> Shaped '[Product sh] a +sflatten arr = + case shsProduct (sshape arr) of -- TODO: simplify when removing the KnownNat stuff + n@SNat -> sreshape (n :$$ ZSS) arr + siota :: (Enum a, PrimElt a) => SNat n -> Shaped '[n] a siota sn = Shaped (miota sn) -- cgit v1.2.3-70-g09d2