aboutsummaryrefslogtreecommitdiff
path: root/src
diff options
context:
space:
mode:
Diffstat (limited to 'src')
-rw-r--r--src/Data/Array/Nested.hs6
-rw-r--r--src/Data/Array/Nested/Internal/Mixed.hs6
-rw-r--r--src/Data/Array/Nested/Internal/Ranked.hs6
-rw-r--r--src/Data/Array/Nested/Internal/Shaped.hs6
4 files changed, 21 insertions, 3 deletions
diff --git a/src/Data/Array/Nested.hs b/src/Data/Array/Nested.hs
index b825691..499fdda 100644
--- a/src/Data/Array/Nested.hs
+++ b/src/Data/Array/Nested.hs
@@ -14,7 +14,7 @@ module Data.Array.Nested (
rfromListLinear, rfromListPrimLinear, rtoListLinear,
rslice, rrev1, rreshape, rflatten, riota,
rminIndexPrim, rmaxIndexPrim, rdot1Inner, rdot,
- rnest, runNest,
+ rnest, runNest, rzip, runzip,
-- ** Lifting orthotope operations to 'Ranked' arrays
rlift, rlift2,
-- ** Conversions
@@ -36,7 +36,7 @@ module Data.Array.Nested (
sfromListLinear, sfromListPrimLinear, stoListLinear,
sslice, srev1, sreshape, sflatten, siota,
sminIndexPrim, smaxIndexPrim, sdot1Inner, sdot,
- snest, sunNest,
+ snest, sunNest, szip, sunzip,
-- ** Lifting orthotope operations to 'Shaped' arrays
slift, slift2,
-- ** Conversions
@@ -58,7 +58,7 @@ module Data.Array.Nested (
mfromListLinear, mfromListPrimLinear, mtoListLinear,
mslice, mrev1, mreshape, mflatten, miota,
mminIndexPrim, mmaxIndexPrim, mdot1Inner, mdot,
- mnest, munNest,
+ mnest, munNest, mzip, munzip,
-- ** Lifting orthotope operations to 'Mixed' arrays
mlift, mlift2,
-- ** Conversions
diff --git a/src/Data/Array/Nested/Internal/Mixed.hs b/src/Data/Array/Nested/Internal/Mixed.hs
index 8d239cf..619d9bc 100644
--- a/src/Data/Array/Nested/Internal/Mixed.hs
+++ b/src/Data/Array/Nested/Internal/Mixed.hs
@@ -790,6 +790,12 @@ mnest ssh arr = M_Nest (fst (shxSplitApp (Proxy @sh') ssh (mshape arr))) arr
munNest :: Mixed sh (Mixed sh' a) -> Mixed (sh ++ sh') a
munNest (M_Nest _ arr) = arr
+mzip :: Mixed sh a -> Mixed sh b -> Mixed sh (a, b)
+mzip = M_Tup2
+
+munzip :: Mixed sh (a, b) -> (Mixed sh a, Mixed sh b)
+munzip (M_Tup2 a b) = (a, b)
+
mrerankP :: forall sh1 sh2 sh a b. (Storable a, Storable b)
=> StaticShX sh -> IShX sh2
-> (Mixed sh1 (Primitive a) -> Mixed sh2 (Primitive b))
diff --git a/src/Data/Array/Nested/Internal/Ranked.hs b/src/Data/Array/Nested/Internal/Ranked.hs
index 9e8a7b2..1966270 100644
--- a/src/Data/Array/Nested/Internal/Ranked.hs
+++ b/src/Data/Array/Nested/Internal/Ranked.hs
@@ -390,6 +390,12 @@ runNest rarr@(Ranked (M_Ranked (M_Nest _ arr)))
| Refl <- lemReplicatePlusApp (rrank rarr) (Proxy @m) (Proxy @(Nothing @Nat))
= Ranked arr
+rzip :: Ranked n a -> Ranked n b -> Ranked n (a, b)
+rzip = coerce mzip
+
+runzip :: Ranked n (a, b) -> (Ranked n a, Ranked n b)
+runzip = coerce munzip
+
rrerankP :: forall n1 n2 n a b. (Storable a, Storable b)
=> SNat n -> IShR n2
-> (Ranked n1 (Primitive a) -> Ranked n2 (Primitive b))
diff --git a/src/Data/Array/Nested/Internal/Shaped.hs b/src/Data/Array/Nested/Internal/Shaped.hs
index 228d800..5d3eba1 100644
--- a/src/Data/Array/Nested/Internal/Shaped.hs
+++ b/src/Data/Array/Nested/Internal/Shaped.hs
@@ -376,6 +376,12 @@ sunNest sarr@(Shaped (M_Shaped (M_Nest _ arr)))
| Refl <- lemMapJustApp (sshape sarr) (Proxy @sh')
= Shaped arr
+szip :: Shaped sh a -> Shaped sh b -> Shaped sh (a, b)
+szip = coerce mzip
+
+sunzip :: Shaped sh (a, b) -> (Shaped sh a, Shaped sh b)
+sunzip = coerce munzip
+
srerankP :: forall sh1 sh2 sh a b. (Storable a, Storable b)
=> ShS sh -> ShS sh2
-> (Shaped sh1 (Primitive a) -> Shaped sh2 (Primitive b))