-
Notifications
You must be signed in to change notification settings - Fork 5
Expand file tree
/
Copy pathIndex.hs
More file actions
93 lines (84 loc) · 2.91 KB
/
Index.hs
File metadata and controls
93 lines (84 loc) · 2.91 KB
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
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
--------------------------------------------------------------------------------
-- |
-- Module : ArrayFire.Index
-- Copyright : David Johnson (c) 2019-2020
-- License : BSD 3
-- Maintainer : David Johnson <djohnson.m@gmail.com>
-- Stability : Experimental
-- Portability : GHC
--
-- Functions for indexing into an 'Array'
--
--------------------------------------------------------------------------------
module ArrayFire.Index where
import ArrayFire.Internal.Index
import ArrayFire.Internal.Types
import ArrayFire.FFI
import ArrayFire.Exception
import Foreign
import System.IO.Unsafe
import Control.Exception
-- | Index into an 'Array' by 'Seq'
index
:: Array a
-- ^ 'Array' argument
-> [Seq]
-- ^ 'Seq' to use for indexing
-> Array a
index (Array fptr) seqs =
unsafePerformIO . mask_ . withForeignPtr fptr $ \ptr -> do
alloca $ \aptr ->
withArray (toAFSeq <$> seqs) $ \sptr -> do
throwAFError =<< af_index aptr ptr n sptr
Array <$> do
newForeignPtr af_release_array_finalizer
=<< peek aptr
where
n = fromIntegral (length seqs)
-- | Lookup an Array by keys along a specified dimension
lookup :: Array a -> Array a -> Int -> Array a
lookup a b n = op2 a b $ \p x y -> af_lookup p x y (fromIntegral n)
-- af_err af_assign_seq( af_array *out, const af_array lhs, const unsigned ndims, const af_seq* const indices, const af_array rhs);
-- | Calculates 'mean' of 'Array' along user-specified dimension.
--
-- @
-- >>> print $ mean 0 ( vector @Int 10 [1..] )
-- @
-- @
-- ArrayFire Array
-- [1 1 1 1]
-- 5.5000
-- @
-- assignSeq :: Array a -> Int -> [Seq] -> Array a -> Array a
-- assignSeq = error "Not implemneted"
-- af_err af_index_gen( af_array *out, const af_array in, const dim_t ndims, const af_index_t* indices);
-- | Calculates 'mean' of 'Array' along user-specified dimension.
--
-- @
-- >>> print $ mean 0 ( vector @Int 10 [1..] )
-- @
-- @
-- ArrayFire Array
-- [1 1 1 1]
-- 5.5000
-- @
-- indexGen :: Array a -> Int -> [Index a] -> Array a -> Array a
-- indexGen = error "Not implemneted"
-- af_err af_assingn_gen( af_array *out, const af_array lhs, const dim_t ndims, const af_index_t* indices, const af_array rhs);
-- | Calculates 'mean' of 'Array' along user-specified dimension.
--
-- @
-- >>> print $ mean 0 ( vector @Int 10 [1..] )
-- @
-- @
-- ArrayFire Array
-- [1 1 1 1]
-- 5.5000
-- @
-- assignGen :: Array a -> Int -> [Index a] -> Array a -> Array a
-- assignGen = error "Not implemneted"
-- af_err af_create_indexers(af_index_t** indexers);
-- af_err af_set_array_indexer(af_index_t* indexer, const af_array idx, const dim_t dim);
-- af_err af_set_seq_indexer(af_index_t* indexer, const af_seq* idx, const dim_t dim, const bool is_batch);
-- af_err af_set_seq_param_indexer(af_index_t* indexer, const double begin, const double end, const double step, const dim_t dim, const bool is_batch);
-- af_err af_release_indexers(af_index_t* indexers);