Readd TestFlat

Author: canyon289

pymc/tests/test_distributions_random.py

@@ -314,6 +314,29 @@ def seeded_numpy_distribution_builder(dist_name: str) -> Callable:
     )
 
 
+class TestFlat(BaseTestDistributionRandom):
+    pymc_dist = pm.Flat
+    pymc_dist_params = {}
+    expected_rv_op_params = {}
+    checks_to_run = [
+        "check_pymc_params_match_rv_op",
+        "check_rv_inferred_size",
+        "check_not_implemented",
+    ]
+
+    def check_rv_inferred_size(self):
+        sizes_to_check = [None, (), 1, (1,), 5, (4, 5), (2, 4, 2)]
+        sizes_expected = [(), (), (1,), (1,), (5,), (4, 5), (2, 4, 2)]
+        for size, expected in zip(sizes_to_check, sizes_expected):
+            pymc_rv = self.pymc_dist.dist(**self.pymc_dist_params, size=size)
+            expected_symbolic = tuple(pymc_rv.shape.eval())
+            assert expected_symbolic == expected
+
+    def check_not_implemented(self):
+        with pytest.raises(NotImplementedError):
+            self.pymc_rv.eval()
+
+
 class TestHalfFlat(BaseTestDistributionRandom):
     pymc_dist = pm.HalfFlat
     pymc_dist_params = {}