diff options
| author | Tom Smeding <tom@tomsmeding.com> | 2024-05-17 13:31:39 +0200 |
|---|---|---|
| committer | Tom Smeding <tom@tomsmeding.com> | 2024-05-17 13:31:39 +0200 |
| commit | 43de4dd2e730273eb04bdf7d0ac62ac1e1422880 (patch) | |
| tree | 7a91b756734222842999d767378c7544aac1e3bf /src/Data/Array/Nested/Internal.hs | |
| parent | 03879dfc925510c13c46664ca3438cc1d4872048 (diff) | |
Diffstat (limited to 'src/Data/Array/Nested/Internal.hs')
| -rw-r--r-- | src/Data/Array/Nested/Internal.hs | 17 |
1 files changed, 11 insertions, 6 deletions
diff --git a/src/Data/Array/Nested/Internal.hs b/src/Data/Array/Nested/Internal.hs index 27e1a30..7bd6565 100644 --- a/src/Data/Array/Nested/Internal.hs +++ b/src/Data/Array/Nested/Internal.hs @@ -594,8 +594,11 @@ mconstant :: forall sh a. (KnownShapeX sh, Storable a, PrimElt a) => IShX sh -> a -> Mixed sh a mconstant sh x = fromPrimitive (mconstantP sh x) -mslice :: (KnownShapeX sh, Elt a) => [(Int, Int)] -> Mixed sh a -> Mixed sh a -mslice ivs = mlift $ \_ -> X.slice ivs +mslice :: (KnownShapeX sh, Elt a) => SNat i -> SNat n -> Mixed (Just (i + n + k) : sh) a -> Mixed (Just n : sh) a +mslice i n = withKnownNat n $ mlift $ \_ -> X.slice i n + +msliceU :: (KnownShapeX sh, Elt a) => Int -> Int -> Mixed (Nothing : sh) a -> Mixed (Nothing : sh) a +msliceU i n = mlift $ \_ -> X.sliceU i n mrev1 :: (KnownShapeX (n : sh), Elt a) => Mixed (n : sh) a -> Mixed (n : sh) a mrev1 = mlift $ \_ -> X.rev1 @@ -1158,8 +1161,10 @@ rconstant :: forall n a. (KnownNat n, Storable a, PrimElt a) => IShR n -> a -> Ranked n a rconstant sh x = coerce fromPrimitive (rconstantP sh x) -rslice :: (KnownNat n, Elt a) => [(Int, Int)] -> Ranked n a -> Ranked n a -rslice ivs = rlift $ \_ -> X.slice ivs +rslice :: forall n a. (KnownNat n, Elt a) => Int -> Int -> Ranked (n + 1) a -> Ranked (n + 1) a +rslice i n + | Refl <- X.lemReplicateSucc @(Nothing @Nat) @n + = rlift $ \_ -> X.sliceU i n rrev1 :: forall n a. (KnownNat n, Elt a) => Ranked (n + 1) a -> Ranked (n + 1) a rrev1 = rlift $ \(Proxy @sh') -> @@ -1436,8 +1441,8 @@ sconstant :: forall sh a. (KnownShape sh, Storable a, PrimElt a) => a -> Shaped sh a sconstant x = coerce fromPrimitive (sconstantP @sh x) -sslice :: (KnownShape sh, Elt a) => [(Int, Int)] -> Shaped sh a -> Shaped sh a -sslice ivs = slift $ \_ -> X.slice ivs +sslice :: (KnownShape sh, Elt a) => SNat i -> SNat n -> Shaped (i + n + k : sh) a -> Shaped (n : sh) a +sslice i n = withKnownNat n $ slift $ \_ -> X.slice i n srev1 :: (KnownNat n, KnownShape sh, Elt a) => Shaped (n : sh) a -> Shaped (n : sh) a srev1 = slift $ \_ -> X.rev1 |