diff options
| author | Tom Smeding <tom@tomsmeding.com> | 2024-05-20 23:00:57 +0200 |
|---|---|---|
| committer | Tom Smeding <tom@tomsmeding.com> | 2024-05-20 23:00:57 +0200 |
| commit | dc9db431fea6bee2aa6533b9df7dee44c002f252 (patch) | |
| tree | 97b5284fda7706cdb203b76b03108358476feab3 | |
| parent | 174484b1e372402681de3f171bd851c9bf55e803 (diff) | |
Some missing operations (rlift2, slift2)
| -rw-r--r-- | src/Data/Array/Nested.hs | 5 | ||||
| -rw-r--r-- | src/Data/Array/Nested/Internal.hs | 14 |
2 files changed, 17 insertions, 2 deletions
diff --git a/src/Data/Array/Nested.hs b/src/Data/Array/Nested.hs index 65f2b18..8d9601e 100644 --- a/src/Data/Array/Nested.hs +++ b/src/Data/Array/Nested.hs @@ -12,7 +12,7 @@ module Data.Array.Nested ( rreplicate, rfromListOuter, rfromList1, rtoListOuter, rtoList1, rslice, rrev1, rreshape, -- ** Lifting orthotope operations to 'Ranked' arrays - rlift, + rlift, rlift2, -- ** Conversions rasXArrayPrim, rfromXArrayPrim, rcastToShaped, @@ -29,7 +29,7 @@ module Data.Array.Nested ( sreplicate, sfromListOuter, sfromList1, stoListOuter, stoList1, sslice, srev1, sreshape, -- ** Lifting orthotope operations to 'Shaped' arrays - slift, + slift, slift2, -- ** Conversions sasXArrayPrim, sfromXArrayPrim, stoRanked, @@ -38,6 +38,7 @@ module Data.Array.Nested ( Mixed, IxX(..), IIxX, KnownShX(..), StaticShX(..), + -- TODO: missing msumOuter1? mshape, mindex, mindexPartial, mgenerate, mtranspose, mappend, mscalar, mfromVector, mtoVector, munScalar, mrerank, 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)) = |