randperm_cpu Class — pytorch Architecture

Source Code

aten/src/ATen/native/TensorFactories.cpp lines 1460–1502

template <typename scalar_t>
void randperm_cpu(Tensor& result, int64_t n, CPUGeneratorImpl* generator) {
  scalar_t* r__data = result.data_ptr<scalar_t>();

  result.resize_({n});
  int64_t r__stride_0 = result.stride(0);

  // for small n, preserve old behavior
  if (n < std::numeric_limits<uint32_t>::max() / 20) {
    at::parallel_for(
        0,
        n,
        internal::GRAIN_SIZE,
        [&r__data, &r__stride_0](int64_t p_begin, int64_t p_end) {
          for (const auto i : c10::irange(p_begin, p_end)) {
            r__data[i * r__stride_0] = static_cast<scalar_t>(i);
          }
        });

    for (int64_t i = 0; i < n - 1; i++) {
      // NOLINTNEXTLINE(clang-analyzer-security.insecureAPI.rand)
      int64_t z = generator->random() % (n - i);
      scalar_t save = r__data[i * r__stride_0];
      r__data[i * r__stride_0] = r__data[(z + i) * r__stride_0];
      r__data[(z + i) * r__stride_0] = save;
    }
    return;
  }

  // we need to pick a number uniformly distributed between 0 and n
  // when n is of the same order of magnitude as the biggest number returned by
  // random the % result is not uniformly distributed
  // so we use random64(), you'd run out of RAM before you
  // start seeing the skew
  // use no-initialization Fischer-Yates variant
  // https://en.wikipedia.org/wiki/Fisher%E2%80%93Yates_shuffle#The_.22inside-out.22_algorithm
  for (int64_t i = 0; i < n; i++) {
    int64_t z = static_cast<int64_t>(generator->random64() % (i + 1));
    r__data[i * r__stride_0] = i;
    r__data[i * r__stride_0] = r__data[z * r__stride_0];
    r__data[z * r__stride_0] = i;
  }
}