xref: /aosp_15_r20/external/pytorch/aten/src/ATen/native/cuda/DistributionGeometricKernel.cu (revision da0073e96a02ea20f0ac840b70461e3646d07c45) 
1 #define TORCH_ASSERT_NO_OPERATORS
2 #include <ATen/cuda/CUDAGeneratorImpl.h>
3 #include <ATen/native/UnaryOps.h>
4 #include <ATen/native/cuda/DistributionTemplates.h>
5 
6 namespace at::native {
7 
geometric_kernel(TensorIteratorBase & iter,double p_,std::optional<Generator> gen)8 void geometric_kernel(TensorIteratorBase& iter, double p_, std::optional<Generator> gen) {
9   auto generator = get_generator_or_default<CUDAGeneratorImpl>(gen, cuda::detail::getDefaultCUDAGenerator());
10   at::native::templates::cuda::geometric_kernel(iter, p_, generator);
11 }
12 
13 REGISTER_DISPATCH(geometric_stub, &geometric_kernel);
14 
15 } // namespace at::native
16