normal_and_transform Class — pytorch Architecture

Source Code

aten/src/ATen/native/cuda/DistributionTemplates.h lines 442–455

template<typename scalar_t, typename accscalar_t, typename RNG, typename transform_t>
void normal_and_transform(TensorIteratorBase& iter, RNG gen, transform_t transform) {
  if (std::is_same_v<scalar_t, double>) {
    distribution_nullary_kernel<scalar_t, accscalar_t, double2>(iter,
      gen,
      [] __device__ (curandStatePhilox4_32_10_t* state) -> double2 { return curand_normal2_double(state); },
      transform);
  } else {
    distribution_nullary_kernel<scalar_t, accscalar_t, float4>(iter,
      gen,
      [] __device__ (curandStatePhilox4_32_10_t* state) -> float4 { return curand_normal4(state); },
      transform);
  }
}