Refactor TruncatedNormal get_moment test

Author: ricardoV94

pymc/distributions/continuous.py

@@ -717,16 +717,16 @@ def dist(
     def get_moment(rv, size, mu, sigma, lower, upper):
         mu, _, lower, upper = at.broadcast_arrays(mu, sigma, lower, upper)
         moment = at.switch(
-            at.isinf(lower),
+            at.eq(lower, -np.inf),
             at.switch(
-                at.isinf(upper),
+                at.eq(upper, np.inf),
                 # lower = -inf, upper = inf
                 mu,
                 # lower = -inf, upper = x
                 upper - 1,
             ),
             at.switch(
-                at.isinf(upper),
+                at.eq(upper, np.inf),
                 # lower = x, upper = inf
                 lower + 1,
                 # lower = x, upper = x

pymc/tests/test_distributions_moments.py

@@ -249,14 +249,13 @@ def test_halfstudentt_moment(nu, sigma, size, expected):
     assert_moment_is_expected(model, expected)
 
 
-@pytest.mark.skip(reason="aeppl interval transform fails when both edges are None")
 @pytest.mark.parametrize(
     "mu, sigma, lower, upper, size, expected",
     [
-        (0.9, 1, -1, 1, None, 0),
-        (0.9, 1, -np.inf, np.inf, 5, np.full(5, 0.9)),
+        (0.9, 1, -5, 5, None, 0),
+        (1, np.ones(5), -10, np.inf, None, np.full(5, -9)),
         (np.arange(5), 1, None, 10, (2, 5), np.full((2, 5), 9)),
-        (1, np.ones(5), -10, np.inf, None, np.full((2, 5), -9)),
+        (1, 1, [-np.inf, -np.inf, -np.inf], 10, None, np.full(3, 9)),
     ],
 )
 def test_truncatednormal_moment(mu, sigma, lower, upper, size, expected):