From 9b0651bf19e889dfb28ba81b6ada25b27b0e6071 Mon Sep 17 00:00:00 2001
From: Tom Smeding <tom@tomsmeding.com>
Date: Mon, 17 Jun 2024 13:08:13 +0200
Subject: sumAllPrim

---
 src/Data/Array/Nested/Internal/Mixed.hs  | 3 +++
 src/Data/Array/Nested/Internal/Ranked.hs | 3 +++
 src/Data/Array/Nested/Internal/Shaped.hs | 3 +++
 3 files changed, 9 insertions(+)

(limited to 'src/Data/Array/Nested')

diff --git a/src/Data/Array/Nested/Internal/Mixed.hs b/src/Data/Array/Nested/Internal/Mixed.hs
index 594383c..215313e 100644
--- a/src/Data/Array/Nested/Internal/Mixed.hs
+++ b/src/Data/Array/Nested/Internal/Mixed.hs
@@ -713,6 +713,9 @@ msumOuter1 :: forall sh n a. (NumElt a, PrimElt a)
            => Mixed (n : sh) a -> Mixed sh a
 msumOuter1 = fromPrimitive . msumOuter1P @sh @n @a . toPrimitive
 
+msumAllPrim :: (PrimElt a, NumElt a) => Mixed sh a -> a
+msumAllPrim (toPrimitive -> M_Primitive sh arr) = X.sumFull (ssxFromShape sh) arr
+
 mappend :: forall n m sh a. Elt a
         => Mixed (n : sh) a -> Mixed (m : sh) a -> Mixed (AddMaybe n m : sh) a
 mappend arr1 arr2 = mlift2 (snm :!% ssh) f arr1 arr2
diff --git a/src/Data/Array/Nested/Internal/Ranked.hs b/src/Data/Array/Nested/Internal/Ranked.hs
index bd37e7a..74b2186 100644
--- a/src/Data/Array/Nested/Internal/Ranked.hs
+++ b/src/Data/Array/Nested/Internal/Ranked.hs
@@ -282,6 +282,9 @@ rsumOuter1 :: forall n a. (NumElt a, PrimElt a)
            => Ranked (n + 1) a -> Ranked n a
 rsumOuter1 = rfromPrimitive . rsumOuter1P . rtoPrimitive
 
+rsumAllPrim :: (PrimElt a, NumElt a) => Ranked n a -> a
+rsumAllPrim (Ranked arr) = msumAllPrim arr
+
 rtranspose :: forall n a. Elt a => PermR -> Ranked n a -> Ranked n a
 rtranspose perm arr
   | sn@SNat <- rrank arr
diff --git a/src/Data/Array/Nested/Internal/Shaped.hs b/src/Data/Array/Nested/Internal/Shaped.hs
index f50ed28..ea979fa 100644
--- a/src/Data/Array/Nested/Internal/Shaped.hs
+++ b/src/Data/Array/Nested/Internal/Shaped.hs
@@ -277,6 +277,9 @@ ssumOuter1 :: forall sh n a. (NumElt a, PrimElt a)
            => Shaped (n : sh) a -> Shaped sh a
 ssumOuter1 = sfromPrimitive . ssumOuter1P . stoPrimitive
 
+ssumAllPrim :: (PrimElt a, NumElt a) => Shaped n a -> a
+ssumAllPrim (Shaped arr) = msumAllPrim arr
+
 stranspose :: forall is sh a. (IsPermutation is, Rank is <= Rank sh, Elt a)
            => Perm is -> Shaped sh a -> Shaped (PermutePrefix is sh) a
 stranspose perm sarr@(Shaped arr)
-- 
cgit v1.2.3-70-g09d2