Fix underflow issue with dirichlet sample (#17488)
authorNeeraj Pradhan <npradhan@uber.com>
Tue, 19 Mar 2019 17:18:12 +0000 (10:18 -0700)
committerFacebook Github Bot <facebook-github-bot@users.noreply.github.com>
Tue, 19 Mar 2019 17:34:13 +0000 (10:34 -0700)
commit9a153412fd4f78b9a9b59bbf85a358339fb69613
treed33e3fcdf582ea7bdb1c33efe42cc2ea8743714c
parent84fe20600dc55c204577784b8902b4af880fab26
Fix underflow issue with dirichlet sample (#17488)

Summary:
Addresses #15738, using fritzo's suggestion. This adds a `torch._sample_dirichlet` method in `Distributions.cpp` and `Distributions.cu`.
 - For CPU, this leads to no perf hit since all we do is to promote the `alpha` to double when getting the gamma samples (the gamma sampler anyways uses `accscalar_t`(double for CPU)) and cast it back to float32 on return.
 - I have added an analogous method for CUDA as well, but the default sampler for CUDA uses scalar_t for efficiency, so I have kept it as that. With this, I do not see the bias towards 1 as reported in #15738 with `float32`, but there is a spurious mode at 0.5, as would be expected. Users would need to explicitly use `float64` for GPU to not see the spurious mode at 0.5. (EDIT: see note below, it appears that the bias issue is still there for certain builds).

Added some tests and checked that there is no perf regression. My experience with C++ is very limited, so apologies in advance if I missed something basic. cc. ailzhang, fritzo, fmassa
Pull Request resolved: https://github.com/pytorch/pytorch/pull/17488

Differential Revision: D14410301

Pulled By: ezyang

fbshipit-source-id: 62b2f694b4642685eab06db96d74ce28e05c3992
aten/src/ATen/native/Distributions.cpp
aten/src/ATen/native/cuda/Distributions.cu
aten/src/ATen/native/native_functions.yaml
test/test_distributions.py
torch/distributions/dirichlet.py