Fix distributions tests

Author: lucianopaz

homepy/blocks/distributions.py

@@ -130,13 +130,8 @@ def rng_fn(cls, rng, dim, alpha, size):
         return v / np.sqrt(np.sum(v**2, axis=-1, keepdims=True) + W[..., None])
 
     def _supp_shape_from_params(self, dist_params, param_shapes=None, rep_param_idx=1):
-        dim = dist_params[0]
-        ref_param = dist_params[rep_param_idx]
-        if param_shapes is not None:
-            ref_param_shape = param_shapes[rep_param_idx]
-            return (*ref_param_shape, dim.flatten()[0])
-        else:
-            return (*ref_param.shape, dim.flatten()[0])
+        dim = dist_params[0].flatten()[0]
+        return (dim,)
 
 
 hyperballUniformRV = HyperballUniformRV()
@@ -156,9 +151,10 @@ def dist(cls, dim, alpha=1.0, no_assert: bool = False, **kwargs):
     @staticmethod
     def support_point(rv, size, dim, alpha):
         """Define the moment (initial point) for the RV"""
+        dim = dim.flatten()[0]
         moment = pt.ones((dim,), dtype=rv.dtype) * 0.5 / pt.sqrt(dim)
         if not rv_size_is_none(size):
-            moment_size = pt.concatenate([size, [dim.flatten()[0]]])
+            moment_size = pt.concatenate([size, [dim]])
             moment = pt.full(moment_size, moment)
         return moment
 
@@ -168,8 +164,17 @@ def ball_transform(op, rv):
     return ballTransform
 
 
+def log1mexp_numpy(x):
+    return np.where(
+        x < np.log(0.5),
+        np.log1p(-np.exp(x)),
+        np.log(-np.expm1(x)),
+    )
+
+
 @_logprob.register(HyperballUniformRV)
 def logp(op, values, rng, size, dim, alpha, **kwargs):
+    dim = dim.flatten()[0]
     value = values[0]
     norm = pt.sqrt(pt.sum(value**2, axis=-1))
     return check_parameters(
@@ -239,7 +244,7 @@ def _inverse_rng_fn(cls, rng, theta, lam, dist_size, idxs_mask):
                 log1p_lam_m_C = np.where(
                     pos_lam,
                     np.log1p(np.exp(abs_log_lam - log_c)),
-                    pm.math.log1mexp_numpy(abs_log_lam - log_c, negative_input=True),
+                    log1mexp_numpy(abs_log_lam - log_c),
                 )
                 log_p = log_c + log1p_lam_m_C * (x_ - 1) + log_p - np.log(x_) - lam
                 log_s = np.logaddexp(log_s, log_p)

homepy/blocks/likelihoods.py

@@ -38,7 +38,10 @@ def get_link(link_function=None):
         try:
             link = getattr(pm.math, link_function)
         except AttributeError:
-            link = getattr(pt, link_function)
+            try:
+                link = getattr(pt, link_function)
+            except AttributeError:
+                link = getattr(pt.special, link_function)
     else:
         raise TypeError(
             f"Cannot understand supplied link function type {type(link_function)}."

homepy/tests/blocks/test_gp.py

@@ -552,7 +552,7 @@ def test_build_correlator_matrix(dims):
     assert m.coords["feature_correlator_"] == (*range(dims - 1), "ignored")
     ball_rv = m["prefix_correlator"]
     assert isinstance(ball_rv.owner.op, HyperballUniformRV)
-    _, size, _, dim, alpha = m["prefix_correlator"].owner.inputs
+    _, size, dim, alpha = m["prefix_correlator"].owner.inputs
     assert tuple(size.eval()) == ()
     assert dim.eval() == dims
     assert alpha.eval() == 1.5

homepy/tests/blocks/test_likelihoods.py

@@ -141,7 +141,7 @@ def test_parse_observed_features_and_dims_error():
         (
             NegBinLikelihood,
             dict(concentration_prior=[1.0]),
-            ("add", "exp", "exponential_rv", "nbinom_rv"),
+            ("add", "exp", "exponential_rv", "negative_binomial_rv"),
             ("Abs", "sqr", "halfnormal_rv"),
         ),
         (
@@ -151,7 +151,7 @@ def test_parse_observed_features_and_dims_error():
                 concentration_dist=pm.Exponential,
                 concentration_prior=dict(lam=pt.sqr(1.0)),
             ),
-            ("add", "exp", "sqr", "exponential_rv", "nbinom_rv"),
+            ("add", "exp", "sqr", "exponential_rv", "negative_binomial_rv"),
             ("Abs", "halfnormal_rv"),
         ),
         (
@@ -162,7 +162,7 @@ def test_parse_observed_features_and_dims_error():
                 concentration_name="concentration",
                 concentration_prior=None,
             ),
-            ("add", "Abs", "halfnormal_rv", "nbinom_rv"),
+            ("add", "Abs", "halfnormal_rv", "negative_binomial_rv"),
             ("exp", "sqr", "exponential_rv"),
         ),
         (