From 29a80650f8d5c6b2df4af96c2b228a8c2f9593bc Mon Sep 17 00:00:00 2001 From: Anatoly Rubanov Date: Mon, 25 Sep 2023 13:19:09 +0300 Subject: [PATCH] change normal ditribution by uniform in step func --- tests/distributions/test_distribution.py | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/tests/distributions/test_distribution.py b/tests/distributions/test_distribution.py index 8cea687f55f..a148c4eeba2 100644 --- a/tests/distributions/test_distribution.py +++ b/tests/distributions/test_distribution.py @@ -473,22 +473,25 @@ def test_custom_dist_default_moment(self, dist_params, size, expected, dist_fn): assert_moment_is_expected(model, expected) def test_custom_dist_default_moment_inner_graph(self): - def scan_step(mu): - x = pm.Normal.dist(mu, 1) + def scan_step(left, right): + x = pm.Uniform.dist(left, right) x_update = collect_default_updates([x]) return x, x_update def dist(size): xs, updates = scan( fn=scan_step, - sequences=[pt.ones(2)], + sequences=[ + pt.as_tensor_variable(np.array([-4, -3])), + pt.as_tensor_variable(np.array([-2, -1])), + ], name="xs", ) - return pt.sum(xs) + return xs with Model() as model: CustomDist("x", dist=dist) - assert_moment_is_expected(model, 2, check_finite_logp=False) + assert_moment_is_expected(model, np.array([-3, -2])) def test_logcdf_inference(self): def custom_dist(mu, sigma, size):