Split distribution operator test cases (#2609)

Author: tengqm

tests/test_bernoulli.py

@@ -0,0 +1,46 @@
+import pytest
+import torch
+
+import flag_gems
+
+from . import accuracy_utils as utils
+
+
+@pytest.mark.bernoulli_
+@pytest.mark.parametrize("shape", utils.DISTRIBUTION_SHAPES)
+@pytest.mark.parametrize("dtype", utils.FLOAT_DTYPES)
+def test_bernoulli_(shape, dtype):
+    x = torch.empty(size=shape, dtype=dtype, device=flag_gems.device)
+    p = 0.5
+    with flag_gems.use_gems():
+        x.bernoulli_(p)
+
+    # Check that all values are 0 or 1
+    assert ((x == 0) | (x == 1)).all()
+
+    # Check that the mean is approximately p (statistical test)
+    mean = x.float().mean().item()
+    assert abs(mean - p) < 0.1
+
+
+@pytest.mark.bernoulli_
+@pytest.mark.parametrize("shape", utils.DISTRIBUTION_SHAPES)
+@pytest.mark.parametrize("dtype", utils.FLOAT_DTYPES)
+@pytest.mark.parametrize("p", [0.0, 0.3, 0.7, 1.0])
+def test_bernoulli_various_p(shape, dtype, p):
+    x = torch.empty(size=shape, dtype=dtype, device=flag_gems.device)
+    with flag_gems.use_gems():
+        x.bernoulli_(p)
+
+    # Check that all values are 0 or 1
+    assert ((x == 0) | (x == 1)).all()
+
+    # Check boundary cases
+    if p == 0.0:
+        assert (x == 0).all()
+    elif p == 1.0:
+        assert (x == 1).all()
+    else:
+        # Check that the mean is approximately p
+        mean = x.float().mean().item()
+        assert abs(mean - p) < 0.15

tests/test_distribution_ops.py

@@ -0,0 +1,38 @@
+import pytest
+import torch
+
+import flag_gems
+
+from . import accuracy_utils as utils
+
+
+@pytest.mark.exponential_
+@pytest.mark.parametrize("shape", utils.DISTRIBUTION_SHAPES)
+@pytest.mark.parametrize("dtype", utils.FLOAT_DTYPES)
+def test_exponential_(shape, dtype):
+    x = torch.empty(size=shape, dtype=dtype, device=flag_gems.device)
+    with flag_gems.use_gems():
+        x.exponential_()
+
+    assert x.min() > 0
+
+
+@pytest.mark.exponential_
+@pytest.mark.parametrize("shape", utils.DISTRIBUTION_SHAPES)
+@pytest.mark.parametrize("dtype", utils.FLOAT_DTYPES)
+def test_accuracy_fast_exponential_(shape, dtype):
+    x = torch.empty(size=shape, dtype=dtype, device=flag_gems.device)
+    lambd = 1.0
+    mean_tol = 0.05
+    var_tol = 0.05
+    with flag_gems.use_gems():
+        x.exponential_()
+
+    x_res = utils.to_reference(x)
+    mean_res = torch.mean(x_res.to(torch.float32)).to(dtype)
+    var_res = torch.var(x_res.to(torch.float32)).to(dtype)
+    mean_ref = 1.0 / lambd
+    var_ref = 1.0 / (lambd**2)
+
+    assert torch.abs(mean_res - mean_ref) < mean_tol
+    assert torch.abs(var_res - var_ref) < var_tol

tests/test_exponential.py

@@ -1,7 +1,9 @@
 import random
 import time
 
+import numpy as np
 import pytest
+import scipy
 import torch
 
 import flag_gems
@@ -43,6 +45,26 @@ def test_multinomial_with_replacement(shape, dtype, n_samples):
             assert torch.sum(res_dist == 0) / res_dist.numel() < 0.001
 
 
+@pytest.mark.multinomial
+@pytest.mark.parametrize("shape", [(1024, 10)])
+@pytest.mark.parametrize("dtype", [torch.float16, torch.float32])
+@pytest.mark.parametrize("n_samples", [2048])
+def test_multinomial_with_replacement_1(shape, dtype, n_samples):
+    # First use multinomial to generate a series of indices, then
+    # use the index counts as the input probabilities (scaled)
+    rand_indices = torch.multinomial(torch.rand(shape), n_samples, True).to(device)
+    inp_counts = torch.nn.functional.one_hot(rand_indices).sum(1)
+    with flag_gems.use_gems():
+        out_indices = torch.multinomial(inp_counts.to(dtype=dtype), n_samples, True)
+    out_counts = torch.nn.functional.one_hot(out_indices).sum(1)
+
+    # Do a simple Chi-square test
+    assert torch.equal(inp_counts.sum(-1), out_counts.sum(-1))
+
+    _, pvalue = scipy.stats.chisquare(out_counts.tolist(), inp_counts.tolist(), axis=-1)
+    assert np.sum(pvalue < 0.05) / len(pvalue) < 0.1
+
+
 @pytest.mark.multinomial
 @pytest.mark.parametrize("pool", utils.UT_SHAPES_2D)
 @pytest.mark.parametrize("dtype", utils.FLOAT_DTYPES)