| 82 | } |
| 83 | |
| 84 | af_err af_index(af_array* result, const af_array in, const unsigned ndims, |
| 85 | const af_seq* indices) { |
| 86 | try { |
| 87 | ARG_ASSERT(2, (ndims > 0 && ndims <= AF_MAX_DIMS)); |
| 88 | |
| 89 | const ArrayInfo& inInfo = getInfo(in); |
| 90 | af_dtype type = inInfo.getType(); |
| 91 | const dim4& iDims = inInfo.dims(); |
| 92 | |
| 93 | vector<af_seq> indices_(ndims, af_span); |
| 94 | for (unsigned i = 0; i < ndims; ++i) { |
| 95 | indices_[i] = convert2Canonical(indices[i], iDims[i]); |
| 96 | |
| 97 | ARG_ASSERT(3, (indices_[i].begin >= 0. && indices_[i].end >= 0.)); |
| 98 | if (signbit(indices_[i].step)) { |
| 99 | ARG_ASSERT(3, indices_[i].begin >= indices_[i].end); |
| 100 | } else { |
| 101 | ARG_ASSERT(3, indices_[i].begin <= indices_[i].end); |
| 102 | } |
| 103 | } |
| 104 | |
| 105 | af_array out = 0; |
| 106 | |
| 107 | switch (type) { |
| 108 | case f32: out = indexBySeqs<float>(in, indices_); break; |
| 109 | case c32: out = indexBySeqs<cfloat>(in, indices_); break; |
| 110 | case f64: out = indexBySeqs<double>(in, indices_); break; |
| 111 | case c64: out = indexBySeqs<cdouble>(in, indices_); break; |
| 112 | case b8: out = indexBySeqs<char>(in, indices_); break; |
| 113 | case s32: out = indexBySeqs<int>(in, indices_); break; |
| 114 | case u32: out = indexBySeqs<unsigned>(in, indices_); break; |
| 115 | case s16: out = indexBySeqs<short>(in, indices_); break; |
| 116 | case u16: out = indexBySeqs<ushort>(in, indices_); break; |
| 117 | case s64: out = indexBySeqs<intl>(in, indices_); break; |
| 118 | case u64: out = indexBySeqs<uintl>(in, indices_); break; |
| 119 | case s8: out = indexBySeqs<schar>(in, indices_); break; |
| 120 | case u8: out = indexBySeqs<uchar>(in, indices_); break; |
| 121 | case f16: out = indexBySeqs<half>(in, indices_); break; |
| 122 | default: TYPE_ERROR(1, type); |
| 123 | } |
| 124 | swap(*result, out); |
| 125 | } |
| 126 | CATCHALL |
| 127 | return AF_SUCCESS; |
| 128 | } |
| 129 | |
| 130 | template<typename T, typename idx_t> |
| 131 | inline af_array lookup(const af_array& in, const af_array& idx, |