Merge branch 'main' into ciguaran_integrate_blackjax_smc

Author: ciguaran

pymc/logprob/rewriting.py

@@ -70,6 +70,7 @@
 from pytensor.tensor.elemwise import DimShuffle, Elemwise
 from pytensor.tensor.random.rewriting import local_subtensor_rv_lift
 from pytensor.tensor.rewriting.basic import register_canonicalize
+from pytensor.tensor.rewriting.math import local_exp_over_1_plus_exp
 from pytensor.tensor.rewriting.shape import ShapeFeature
 from pytensor.tensor.rewriting.uncanonicalize import local_max_and_argmax
 from pytensor.tensor.subtensor import (
@@ -359,7 +360,12 @@ def incsubtensor_rv_replace(fgraph, node):
 
 logprob_rewrites_db = SequenceDB()
 logprob_rewrites_db.name = "logprob_rewrites_db"
+# Introduce sigmoid. We do it before canonicalization so that useless mul are removed next
+logprob_rewrites_db.register(
+    "local_exp_over_1_plus_exp", out2in(local_exp_over_1_plus_exp), "basic"
+)
 logprob_rewrites_db.register("pre-canonicalize", optdb.query("+canonicalize"), "basic")
+# Split max_and_argmax
 logprob_rewrites_db.register("local_max_and_argmax", out2in(local_max_and_argmax), "basic")
 
 # These rewrites convert un-measurable variables into their measurable forms,

pymc/math.py

@@ -32,13 +32,21 @@
 from pytensor.tensor import (
     abs,
     and_,
+    arccos,
+    arccosh,
+    arcsin,
+    arcsinh,
+    arctan,
+    arctanh,
     broadcast_to,
     ceil,
     clip,
     concatenate,
     constant,
     cos,
     cosh,
+    cumprod,
+    cumsum,
     dot,
     eq,
     erf,
@@ -59,8 +67,10 @@
     logsumexp,
     lt,
     matmul,
+    max,
     maximum,
     mean,
+    min,
     minimum,
     neq,
     ones,
@@ -94,13 +104,21 @@
 __all__ = [
     "abs",
     "and_",
+    "arccos",
+    "arccosh",
+    "arcsin",
+    "arcsinh",
+    "arctan",
+    "arctanh",
     "broadcast_to",
     "ceil",
     "clip",
     "concatenate",
     "constant",
     "cos",
     "cosh",
+    "cumprod",
+    "cumsum",
     "dot",
     "eq",
     "erf",
@@ -121,14 +139,17 @@
     "logsumexp",
     "lt",
     "matmul",
+    "max",
     "maximum",
     "mean",
+    "min",
     "minimum",
     "neq",
     "ones",
     "ones_like",
     "or_",
     "prod",
+    "round",
     "sgn",
     "sigmoid",
     "sin",
@@ -258,7 +279,7 @@ def kron_diag(*diags):
     return reduce(flat_outer, diags)
 
 
-def tround(*args, **kwargs):
+def round(*args, **kwargs):
     """
     Temporary function to silence round warning in PyTensor. Please remove
     when the warning disappears.
@@ -267,6 +288,11 @@ def tround(*args, **kwargs):
     return pt.round(*args, **kwargs)
 
 
+def tround(*args, **kwargs):
+    warnings.warn("tround is deprecated. Use round instead.")
+    return round(*args, **kwargs)
+
+
 def logdiffexp(a, b):
     """log(exp(a) - exp(b))"""
     return a + pt.log1mexp(b - a)

pymc/model/core.py

@@ -1347,15 +1347,6 @@ def make_obs_var(
                 "Dimensionality of data and RV don't match.", actual=data.ndim, expected=rv_var.ndim
             )
 
-        if pytensor.config.compute_test_value != "off":
-            test_value = getattr(rv_var.tag, "test_value", None)
-
-            if test_value is not None:
-                # We try to reuse the old test value
-                rv_var.tag.test_value = np.broadcast_to(test_value, rv_var.shape)
-            else:
-                rv_var.tag.test_value = data
-
         mask = getattr(data, "mask", None)
         if mask is not None:
             impute_message = (

tests/logprob/test_transforms.py

@@ -1127,31 +1127,31 @@ def test_cosh_rv_transform():
     )
 
 
-TRANSFORMATIONS = {
-    "log1p": (pt.log1p, lambda x: pt.log(1 + x)),
-    "softplus": (pt.softplus, lambda x: pt.log(1 + pt.exp(x))),
-    "log1mexp": (pt.log1mexp, lambda x: pt.log(1 - pt.exp(x))),
-    "log2": (pt.log2, lambda x: pt.log(x) / pt.log(2)),
-    "log10": (pt.log10, lambda x: pt.log(x) / pt.log(10)),
-    "exp2": (pt.exp2, lambda x: pt.exp(pt.log(2) * x)),
-    "expm1": (pt.expm1, lambda x: pt.exp(x) - 1),
-    "sigmoid": (pt.sigmoid, lambda x: 1 / (1 + pt.exp(-x))),
-}
-
-
-@pytest.mark.parametrize("transform", TRANSFORMATIONS.keys())
-def test_special_log_exp_transforms(transform):
+@pytest.mark.parametrize(
+    "canonical_func,raw_func",
+    [
+        (pt.log1p, lambda x: pt.log(1 + x)),
+        (pt.softplus, lambda x: pt.log(1 + pt.exp(x))),
+        (pt.log1mexp, lambda x: pt.log(1 - pt.exp(x))),
+        (pt.log2, lambda x: pt.log(x) / pt.log(2)),
+        (pt.log10, lambda x: pt.log(x) / pt.log(10)),
+        (pt.exp2, lambda x: pt.exp(pt.log(2) * x)),
+        (pt.expm1, lambda x: pt.exp(x) - 1),
+        (pt.sigmoid, lambda x: 1 / (1 + pt.exp(-x))),
+        (pt.sigmoid, lambda x: pt.exp(x) / (1 + pt.exp(x))),
+    ],
+)
+def test_special_log_exp_transforms(canonical_func, raw_func):
     base_rv = pt.random.normal(name="base_rv")
     vv = pt.scalar("vv")
 
-    transform_func, ref_func = TRANSFORMATIONS[transform]
-    transformed_rv = transform_func(base_rv)
-    ref_transformed_rv = ref_func(base_rv)
+    transformed_rv = raw_func(base_rv)
+    ref_transformed_rv = canonical_func(base_rv)
 
     logp_test = logp(transformed_rv, vv)
     logp_ref = logp(ref_transformed_rv, vv)
 
-    if transform in ["log2", "log10"]:
+    if canonical_func in (pt.log2, pt.log10):
         # in the cases of log2 and log10 floating point inprecision causes failure
         # from equal_computations so evaluate logp and check all close instead
         vv_test = np.array(0.25)

tests/model/test_core.py

@@ -223,6 +223,14 @@ def test_observed_type(self):
         assert x1.type.dtype == X.type.dtype
         assert x2.type.dtype == X.type.dtype
 
+    @pytensor.config.change_flags(compute_test_value="raise")
+    def test_observed_compute_test_value(self):
+        data = np.zeros(100)
+        with pm.Model():
+            obs = pm.Normal("obs", mu=pt.zeros_like(data), sigma=1, observed=data)
+        assert obs.tag.test_value.shape == data.shape
+        assert obs.tag.test_value.dtype == data.dtype
+
 
 def test_duplicate_vars():
     with pytest.raises(ValueError) as err: