Add tests for samples

scikit-hep · Oct 20, 2020 · 2490342 · 2490342
1 parent 55b18c3
commit 2490342
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/tests/test_probability.py b/tests/test_probability.py
@@ -20,6 +20,9 @@ def test_poisson(backend):
     )
     assert result.shape == (1, 2)
 
+    sample = probability.Poisson(tb.astensor([10.0, 10.0])).sample((10,))
+    assert sample.shape == (10, 2)
+
 
 def test_normal(backend):
     tb, _ = backend
@@ -43,6 +46,9 @@ def test_normal(backend):
     ).log_prob(tb.astensor([[2.0, 3.0]]))
     assert result.shape == (1, 2)
 
+    sample = probability.Poisson(tb.astensor([10.0, 10.0])).sample((10,))
+    assert sample.shape == (10, 2)
+
 
 def test_joint(backend):
     tb, _ = backend