blob: b1c7031880df8df434e4c38aad624754e80c8c5b (
plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
|
{-# LANGUAGE ImportQualifiedPost #-}
module Data.Array.Mixed.Internal.Arith (
module Data.Array.Mixed.Internal.Arith,
module Data.Array.Strided.Arith,
) where
import Data.Array.Internal qualified as OI
import Data.Array.Internal.RankedG qualified as RG
import Data.Array.Internal.RankedS qualified as RS
import Data.Array.Strided qualified as AS
import Data.Array.Strided.Arith
-- for liftVEltwise1
import Foreign.Storable
import GHC.TypeLits
import Data.Vector.Storable qualified as VS
import Data.Array.Strided.Arith.Internal (stridesDense)
fromO :: RS.Array n a -> AS.Array n a
fromO (RS.A (RG.A sh (OI.T strides offset vec))) = AS.Array sh strides offset vec
toO :: AS.Array n a -> RS.Array n a
toO (AS.Array sh strides offset vec) = RS.A (RG.A sh (OI.T strides offset vec))
liftO1 :: (AS.Array n a -> AS.Array n' b)
-> RS.Array n a -> RS.Array n' b
liftO1 f = toO . f . fromO
liftO2 :: (AS.Array n a -> AS.Array n1 b -> AS.Array n2 c)
-> RS.Array n a -> RS.Array n1 b -> RS.Array n2 c
liftO2 f x y = toO (f (fromO x) (fromO y))
liftVEltwise1 :: (Storable a, Storable b)
=> SNat n
-> (VS.Vector a -> VS.Vector b)
-> RS.Array n a -> RS.Array n b
liftVEltwise1 SNat f arr@(RS.A (RG.A sh (OI.T strides offset vec)))
| Just (blockOff, blockSz) <- stridesDense sh offset strides =
let vec' = f (VS.slice blockOff blockSz vec)
in RS.A (RG.A sh (OI.T strides (offset - blockOff) vec'))
| otherwise = RS.fromVector sh (f (RS.toVector arr))
|
