| 220 | } |
| 221 | |
| 222 | af_err af_index_gen(af_array* out, const af_array in, const dim_t ndims, |
| 223 | const af_index_t* indexs) { |
| 224 | try { |
| 225 | ARG_ASSERT(2, (ndims > 0 && ndims <= AF_MAX_DIMS)); |
| 226 | ARG_ASSERT(3, (indexs != NULL)); |
| 227 | |
| 228 | const ArrayInfo& iInfo = getInfo(in); |
| 229 | const dim4& iDims = iInfo.dims(); |
| 230 | af_dtype inType = getInfo(in).getType(); |
| 231 | |
| 232 | if (iDims.ndims() <= 0) { |
| 233 | *out = createHandle(dim4(0), inType); |
| 234 | return AF_SUCCESS; |
| 235 | } |
| 236 | |
| 237 | if (ndims == 1 && ndims != static_cast<dim_t>(iInfo.ndims())) { |
| 238 | af_array in_ = 0; |
| 239 | AF_CHECK(af_flat(&in_, in)); |
| 240 | AF_CHECK(af_index_gen(out, in_, ndims, indexs)); |
| 241 | AF_CHECK(af_release_array(in_)); |
| 242 | return AF_SUCCESS; |
| 243 | } |
| 244 | |
| 245 | int track = 0; |
| 246 | std::array<af_seq, AF_MAX_DIMS> seqs{}; |
| 247 | seqs.fill(af_span); |
| 248 | for (dim_t i = 0; i < ndims; i++) { |
| 249 | if (indexs[i].isSeq) { |
| 250 | track++; |
| 251 | seqs[i] = indexs[i].idx.seq; |
| 252 | } |
| 253 | } |
| 254 | |
| 255 | if (track == static_cast<int>(ndims)) { |
| 256 | return af_index(out, in, ndims, seqs.data()); |
| 257 | } |
| 258 | |
| 259 | std::array<af_index_t, AF_MAX_DIMS> idxrs{}; |
| 260 | |
| 261 | for (dim_t i = 0; i < AF_MAX_DIMS; ++i) { |
| 262 | if (i < ndims) { |
| 263 | bool isSeq = indexs[i].isSeq; |
| 264 | if (!isSeq) { |
| 265 | // check if all af_arrays have atleast one value |
| 266 | // to enable indexing along that dimension |
| 267 | const ArrayInfo& idxInfo = getInfo(indexs[i].idx.arr); |
| 268 | af_dtype idxType = idxInfo.getType(); |
| 269 | |
| 270 | ARG_ASSERT(3, (idxType != c32)); |
| 271 | ARG_ASSERT(3, (idxType != c64)); |
| 272 | ARG_ASSERT(3, (idxType != b8)); |
| 273 | |
| 274 | idxrs[i] = {{indexs[i].idx.arr}, isSeq, indexs[i].isBatch}; |
| 275 | } else { |
| 276 | // copy the af_seq to local variable |
| 277 | af_seq inSeq = |
| 278 | convert2Canonical(indexs[i].idx.seq, iDims[i]); |
| 279 | ARG_ASSERT(3, (inSeq.begin >= 0. || inSeq.end >= 0.)); |
no test coverage detected