Function ShuffleBatchKernel

paddle/phi/kernels/cpu/shuffle_batch_kernel.cc:22–105

template <typename T, typename Context>
void ShuffleBatchKernel(const Context& dev_ctx,
                        const DenseTensor& x,
                        const DenseTensor& seed,
                        int startup_seed,
                        DenseTensor* out,
                        DenseTensor* shuffleidx,
                        DenseTensor* seed_out) {
  auto x_embed_size = x.dims()[x.dims().size() - 1];
  int elem_size = 1;
  for (auto i = 0; i < x.dims().size() - 1; i++)
    elem_size *= static_cast<int>(x.dims()[i]);

  std::vector<int64_t> idx_vec;  // record shuffled order
  idx_vec.reserve(elem_size);
  for (int i = 0; i < elem_size; i++) {
    idx_vec.push_back(i);
  }
  int64_t seed_int = 0;
  if (seed.initialized()) {
    seed_int = *seed.data<int64_t>();
  } else {
    seed_int = startup_seed;
  }
  std::default_random_engine engine;
  engine.seed(seed_int);

  auto custom_random_shuffle = [&idx_vec]() {
    std::random_device rnd;
    int64_t seed_tmp = rnd();
    std::default_random_engine rng(seed_tmp);
    const int n = static_cast<int>(idx_vec.size());
    std::vector<int> v(n);
    std::iota(v.begin(), v.end(), 0);
    std::vector<bool> visit(n, false);
    while (!v.empty()) {
      std::shuffle(v.begin(), v.end(), rng);
      int tmp = v.back();
      v.pop_back();
      if (v.empty()) {
        std::uniform_int_distribution<int> distr(0, n - 2);
        idx_vec[tmp] = tmp;
        std::swap(idx_vec[tmp], idx_vec[(distr(rng) + tmp + 1) % n]);
        return;
      }
      visit[tmp] = true;
      std::shuffle(v.begin(), v.end(), rng);
      int curr = v.back();
      v.pop_back();
      v.push_back(tmp);
      idx_vec[tmp] = curr;
      while (!visit[curr]) {
        visit[curr] = true;
        std::shuffle(v.begin(), v.end(), rng);
        idx_vec[curr] = v.back();
        v.pop_back();
        curr = static_cast<int>(idx_vec[curr]);
      }
    }