random_full_64_bits_range_kernel Class — pytorch Architecture

Source Code

aten/src/ATen/native/cpu/DistributionTemplates.h lines 39–55

template<typename RNG>
void random_full_64_bits_range_kernel(TensorIteratorBase& iter, RNG generator) {
  AT_DISPATCH_ALL_TYPES_AND(at::ScalarType::BFloat16, iter.dtype(), "random_full_64_bits_range_kernel_cpu", [&] {
    if constexpr (std::is_same_v<scalar_t, int64_t> ||
        std::is_same_v<scalar_t, double> ||
        std::is_same_v<scalar_t, float> ||
        std::is_same_v<scalar_t, at::BFloat16>) {
      std::lock_guard<std::mutex> lock(generator->mutex_);
      cpu_serial_kernel(iter, [generator]() -> scalar_t {
        uniform_int_full_range_distribution<scalar_t> random;
        return random(generator);
      });
    } else {
      TORCH_CHECK(false, "random_full_64_bits_range_kernel_cpu handles only int64, double, float and bfloat16");
    }
  });
}