From dc9db431fea6bee2aa6533b9df7dee44c002f252 Mon Sep 17 00:00:00 2001 From: Tom Smeding Date: Mon, 20 May 2024 23:00:57 +0200 Subject: Some missing operations (rlift2, slift2) --- src/Data/Array/Nested/Internal.hs | 14 ++++++++++++++ 1 file changed, 14 insertions(+) (limited to 'src/Data/Array/Nested/Internal.hs') diff --git a/src/Data/Array/Nested/Internal.hs b/src/Data/Array/Nested/Internal.hs index ab67dcc..e402f1e 100644 --- a/src/Data/Array/Nested/Internal.hs +++ b/src/Data/Array/Nested/Internal.hs @@ -1348,6 +1348,13 @@ rlift :: forall n1 n2 a. Elt a -> Ranked n1 a -> Ranked n2 a rlift sn2 f (Ranked arr) = Ranked (mlift (ssxFromSNat sn2) f arr) +-- | See the documentation of 'mlift2'. +rlift2 :: forall n1 n2 n3 a. Elt a + => SNat n3 + -> (forall sh' b. Storable b => StaticShX sh' -> XArray (Replicate n1 Nothing ++ sh') b -> XArray (Replicate n2 Nothing ++ sh') b -> XArray (Replicate n3 Nothing ++ sh') b) + -> Ranked n1 a -> Ranked n2 a -> Ranked n3 a +rlift2 sn3 f (Ranked arr1) (Ranked arr2) = Ranked (mlift2 (ssxFromSNat sn3) f arr1 arr2) + rsumOuter1P :: forall n a. (Storable a, Num a) => Ranked (n + 1) (Primitive a) -> Ranked n (Primitive a) @@ -1593,6 +1600,13 @@ slift :: forall sh1 sh2 a. Elt a -> Shaped sh1 a -> Shaped sh2 a slift sh2 f (Shaped arr) = Shaped (mlift (X.staticShapeFrom (shCvtSX sh2)) f arr) +-- | See the documentation of 'mlift'. +slift2 :: forall sh1 sh2 sh3 a. Elt a + => ShS sh3 + -> (forall sh' b. Storable b => StaticShX sh' -> XArray (MapJust sh1 ++ sh') b -> XArray (MapJust sh2 ++ sh') b -> XArray (MapJust sh3 ++ sh') b) + -> Shaped sh1 a -> Shaped sh2 a -> Shaped sh3 a +slift2 sh3 f (Shaped arr1) (Shaped arr2) = Shaped (mlift2 (X.staticShapeFrom (shCvtSX sh3)) f arr1 arr2) + ssumOuter1P :: forall sh n a. (Storable a, Num a) => Shaped (n : sh) (Primitive a) -> Shaped sh (Primitive a) ssumOuter1P (Shaped (M_Primitive (SKnown sn :$% sh) arr)) = -- cgit v1.2.3-70-g09d2