aboutsummaryrefslogtreecommitdiff
path: root/src/Data/Array/Nested/Internal/Shaped.hs
diff options
context:
space:
mode:
authorTom Smeding <tom@tomsmeding.com>2024-06-10 23:28:11 +0200
committerTom Smeding <tom@tomsmeding.com>2024-06-10 23:28:11 +0200
commit5f6a64660b16d8f188caca5216e55debf4264611 (patch)
tree7e378c5929126db6c583862220e4163de7b2b3df /src/Data/Array/Nested/Internal/Shaped.hs
parent87484b9adcbaa6b380ed3ba1a499bd227a8863a8 (diff)
Add *flatten
Diffstat (limited to 'src/Data/Array/Nested/Internal/Shaped.hs')
-rw-r--r--src/Data/Array/Nested/Internal/Shaped.hs5
1 files changed, 5 insertions, 0 deletions
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)