array_nonzero.cc 683 Bytes
Newer Older
1
/**
2
 *  Copyright (c) 2020 by Contributors
3
4
 * @file array/cpu/array_nonzero.cc
 * @brief Array nonzero CPU implementation
5
6
7
8
9
10
11
12
 */
#include <dgl/array.h>

namespace dgl {
using runtime::NDArray;
namespace aten {
namespace impl {

13
template <DGLDeviceType XPU, typename IdType>
14
15
16
17
IdArray NonZero(IdArray array) {
  std::vector<int64_t> ret;
  const IdType* data = array.Ptr<IdType>();
  for (int64_t i = 0; i < array->shape[0]; ++i)
18
    if (data[i] != 0) ret.push_back(i);
19
20
21
  return NDArray::FromVector(ret, array->ctx);
}

22
23
template IdArray NonZero<kDGLCPU, int32_t>(IdArray);
template IdArray NonZero<kDGLCPU, int64_t>(IdArray);
24
25
26
27

}  // namespace impl
}  // namespace aten
}  // namespace dgl