Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Add Half support for poisson and use float for Half cumulative distribution on CPU #112124

Closed
wants to merge 2 commits into from
Closed

Add Half support for poisson and use float for Half cumulative distribution on CPU #112124

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
0555666
add Half support for poisson and use float for Half cumulative distri…
CaoE Oct 26, 2023
f91eb43
add Half and BFloat16 tests for poisson
CaoE Oct 30, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
2 changes: 1 addition & 1 deletion aten/src/ATen/native/Distributions.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -470,7 +470,7 @@ Tensor _s_poisson_cpu(const Tensor& lambda, c10::optional<Generator> gen) {
.add_output(ret)
.add_input(lambda)
.build();
AT_DISPATCH_FLOATING_TYPES_AND(at::ScalarType::BFloat16, ret.scalar_type(), "poisson_cpu", [&] {
AT_DISPATCH_FLOATING_TYPES_AND2(at::ScalarType::BFloat16, at::ScalarType::Half, ret.scalar_type(), "poisson_cpu", [&] {
CPUGeneratorImpl* generator = get_generator_or_default<CPUGeneratorImpl>(gen, detail::getDefaultCPUGenerator());
// See Note [Acquire lock when using random generators]
std::lock_guard<std::mutex> lock(generator->mutex_);
Expand Down
10 changes: 6 additions & 4 deletions aten/src/ATen/native/cpu/MultinomialKernel.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,8 @@ namespace at::native {
namespace {

template <typename scalar_t>
void multinomial_with_replacement_apply(
typename std::enable_if_t<!is_reduced_floating_point_v<scalar_t>, void>
multinomial_with_replacement_apply(
Tensor& result,
const Tensor& self,
const int64_t n_sample,
Expand Down Expand Up @@ -121,8 +122,9 @@ void multinomial_with_replacement_apply(
}
}

template <>
void multinomial_with_replacement_apply<BFloat16>(
template <typename scalar_t>
typename std::enable_if_t<is_reduced_floating_point_v<scalar_t>, void>
multinomial_with_replacement_apply(
Tensor& result,
const Tensor& self,
const int64_t n_sample,
Expand All @@ -138,7 +140,7 @@ void multinomial_with_replacement_apply<BFloat16>(
/* cumulative probability distribution vector */
Tensor cum_dist = at::empty({n_categories}, self.options().dtype(kFloat));

const BFloat16* const self_ptr = self.data_ptr<BFloat16>();
const scalar_t* const self_ptr = self.data_ptr<scalar_t>();
float* const cum_dist_ptr = cum_dist.data_ptr<float>();
int64_t* const result_ptr = result.data_ptr<int64_t>();

Expand Down
15 changes: 10 additions & 5 deletions test/distributions/test_distributions.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -891,6 +891,7 @@ def _check_sampler_discrete(self, torch_dist, ref_dist, message,
num_samples=10000, failure_rate=1e-3):
"""Runs a Chi2-test for the support, but ignores tail instead of combining"""
torch_samples = torch_dist.sample((num_samples,)).squeeze()
torch_samples = torch_samples.float() if torch_samples.dtype == torch.bfloat16 else torch_samples
torch_samples = torch_samples.cpu().numpy()
unique, counts = np.unique(torch_samples, return_counts=True)
pmf = ref_dist.pmf(unique)
Expand Down Expand Up @@ -1463,11 +1464,15 @@ def ref_log_prob(ref_rate, idx, x, log_prob):
@unittest.skipIf(not TEST_NUMPY, "Numpy not found")
def test_poisson_sample(self):
set_rng_seed(1) # see Note [Randomized statistical tests]
for rate in [0.1, 1.0, 5.0]:
self._check_sampler_discrete(Poisson(rate),
scipy.stats.poisson(rate),
f'Poisson(lambda={rate})',
failure_rate=1e-3)
saved_dtype = torch.get_default_dtype()
for dtype in [torch.float, torch.double, torch.bfloat16, torch.half]:
torch.set_default_dtype(dtype)
for rate in [0.1, 1.0, 5.0]:
self._check_sampler_discrete(Poisson(rate),
scipy.stats.poisson(rate),
f'Poisson(lambda={rate})',
failure_rate=1e-3)
torch.set_default_dtype(saved_dtype)

@unittest.skipIf(not TEST_CUDA, "CUDA not found")
@unittest.skipIf(not TEST_NUMPY, "Numpy not found")
Expand Down