Add bfloat16 CUDA support to binomial distribution (pytorch#116932)

Now all distributions support bfloat16 as input.

Pull Request resolved: pytorch#116932
Approved by: https://github.com/malfet

Author: Skylion007

aten/src/ATen/native/cuda/Distributions.cu

@@ -167,7 +167,7 @@ void launch_binomial_cuda_kernel(
     std::lock_guard<std::mutex> lock(gen->mutex_);
     rng_engine_inputs = gen->philox_cuda_state(42);
   }
-  AT_DISPATCH_FLOATING_TYPES_AND_HALF(iter.input_dtype(), "binomial_cuda", [&] {
+  AT_DISPATCH_FLOATING_TYPES_AND2(at::ScalarType::Half, at::ScalarType::BFloat16, iter.input_dtype(), "binomial_cuda", [&] {
     binomial_cuda_kernel<scalar_t>(iter, rng_engine_inputs);
   });
 }

test/distributions/test_distributions.py

@@ -1151,6 +1151,9 @@ def test_binomial(self):
             self._gradcheck_log_prob(lambda p: Binomial(total_count, None, p.log()), [p])
         self.assertRaises(NotImplementedError, Binomial(10, p).rsample)
 
+    test_binomial_half = set_default_dtype(torch.float16)(test_binomial)
+    test_binomial_bfloat16 = set_default_dtype(torch.bfloat16)(test_binomial)
+
     @unittest.skipIf(not TEST_NUMPY, "NumPy not found")
     def test_binomial_sample(self):
         set_rng_seed(0)  # see Note [Randomized statistical tests]