assert lambda >= 0 in poisson distribution cuda kernel (pytorch#85906)

fix pytorch#85731 Pull Request resolved: pytorch#85906 Approved by: https://github.com/ngimel
zmmwl · Sep 29, 2022 · c036fb3 · c036fb3
1 parent bc57306
commit c036fb3
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 0 deletions.
diff --git a/aten/src/ATen/native/cuda/Distributions.cu b/aten/src/ATen/native/cuda/Distributions.cu
@@ -47,6 +47,7 @@ void poisson_cuda_kernel(
     at::PhiloxCudaState philox_args) {
   auto functor = [philox_args] __device__(
           scalar_t & ret_val, const scalar_t& lambda) {
+        CUDA_KERNEL_ASSERT(lambda >= 0 && "invalid Poisson rate, expected rate to be non-negative");
         auto seeds = at::cuda::philox::unpack(philox_args);
         curandStatePhilox4_32_10_t state;
         curand_init(std::get<0>(seeds),

diff --git a/torch/_torch_docs.py b/torch/_torch_docs.py
@@ -8514,6 +8514,8 @@ def merge_dicts(*dicts):
 .. math::
     \text{{out}}_i \sim \text{{Poisson}}(\text{{input}}_i)
 
+:attr:`input` must be non-negative.
+
 Args:
     input (Tensor): the input tensor containing the rates of the Poisson distribution