Infer logprob of casts that discretize a continuous variable

A float -> int cast rounds towards zero, so it is the `trunc` a user could
have written themselves composed with a cast that only relabels the dtype.
Rewriting it to that form leaves the judgement of whether the base may be
truncated to `find_measurable_roundings`, and the relabelling cast is then
claimed as measure-preserving.

The other narrowing casts are not truncations and stay unclaimed: casting
to an unsigned int wraps around for negative values (-2.7 -> 254 for
uint8), and casting to bool tests `x != 0`, which collapses the support
onto two points rather than partitioning it.

`round_logprob` snaps the value onto the cell whose mass it asks for, and
the Op it uses to do so rejects the integer value that a cast to int hands
it. The snapping is an identity on such a value, so only its dtype needs
adjusting.

Author: ricardoV94

pymc/logprob/censoring.py

@@ -365,6 +365,13 @@ def round_logprob(op, values, base_rv, **kwargs):
     """
     (value,) = values
 
+    if not value.type.dtype.startswith("float"):
+        # The Ops below snap the value onto the cell whose mass it asks for, and only
+        # accept float inputs. An integral value already sits on a cell, so the snapping
+        # is an identity and only the dtype needs adjusting (as when the rounding is
+        # followed by a cast to int).
+        value = pt.cast(value, "float64")
+
     if isinstance(op.scalar_op, RoundHalfToEven | RoundHalfAwayFromZero):
         # The tie-breaking rule only matters on a measure-zero set of the
         # continuous base variable, so both variants share the same intervals

pymc/logprob/tensor.py

@@ -67,6 +67,7 @@
     _logprob_helper,
     promised_valued_rv,
 )
+from pymc.logprob.censoring import MeasurableRound
 from pymc.logprob.rewriting import (
     assume_valued_outputs,
     early_measurable_ir_rewrites_db,
@@ -487,16 +488,33 @@ def find_measurable_casts(fgraph, node) -> list[TensorVariable] | None:
 
     [out] = node.outputs
     # bool < integer < float; casting to a lower kind discretizes the base variable,
-    # which is not a measure-preserving identity (float -> int would need to be
-    # treated as a form of censoring)
+    # which is not a measure-preserving identity
     kind_order = {"b": 0, "u": 1, "i": 1, "f": 2}
-    in_kind = kind_order.get(np.dtype(base_var.type.dtype).kind)
-    out_kind = kind_order.get(np.dtype(out.type.dtype).kind)
+    in_dtype = np.dtype(base_var.type.dtype)
+    out_dtype = np.dtype(out.type.dtype)
+    in_kind = kind_order.get(in_dtype.kind)
+    out_kind = kind_order.get(out_dtype.kind)
     if in_kind is None or out_kind is None:
         # Kinds we can't order, such as complex
         return None
     if out_kind < in_kind:
-        return None
+        # A float -> signed int cast rounds towards zero, so it is a `trunc` composed
+        # with a cast that merely relabels the dtype. Introduce the `trunc` explicitly
+        # and let `find_measurable_roundings` claim it; the relabelling cast is then
+        # measure-preserving and claimed on a later pass.
+        # The other narrowing casts are not truncations: unsigned ints wrap around for
+        # negative values (-2.7 -> 254 for uint8), and bool tests `x != 0`, which
+        # collapses the support onto two points instead of partitioning it.
+        if in_dtype.kind != "f" or out_dtype.kind != "i":
+            return None
+        # Rewriting to the `trunc` the user could have written themselves leaves the
+        # judgement of whether the base may be truncated to `find_measurable_roundings`,
+        # which declines the bases whose support it cannot establish, rather than
+        # duplicating that reasoning here. Once the `trunc` is in place the cast only
+        # relabels the dtype and is claimed as measure-preserving below; skipping it for
+        # an already rounded base is what brings the rewrite to a fixpoint.
+        if not isinstance(base_var.owner.op, MeasurableRound):
+            return [pt.cast(pt.trunc(base_var), out_dtype.name)]
 
     if in_kind < kind_order["f"] and out_kind == kind_order["f"]:
         # Casting a discrete variable to float hides its discreteness from other

tests/logprob/test_tensor.py

@@ -712,8 +712,38 @@ def test_bool_to_int(self):
             st.bernoulli(0.3).logpmf(1),
         )
 
-    @pytest.mark.parametrize("out_dtype", ["int64", "bool"])
-    def test_discretizing_cast_not_measurable(self, out_dtype):
+    @pytest.mark.parametrize(
+        "value, lower, upper",
+        [(1, 1, 2), (0, -1, 1), (-1, -2, -1)],
+        ids=["positive", "zero", "negative"],
+    )
+    def test_float_to_int(self, value, lower, upper):
+        # The cast rounds towards zero, pooling (-1, 1) at zero
+        y = pt.cast(pt.random.normal(0.5, 1), "int64")
+        y_vv = y.clone()
+
+        np.testing.assert_allclose(
+            logp(y, y_vv).eval({y_vv: value}),
+            np.log(st.norm(0.5, 1).cdf(upper) - st.norm(0.5, 1).cdf(lower)),
+        )
+
+    @pytest.mark.parametrize("rounding_fn", [pt.trunc, pt.floor, pt.ceil, pt.round])
+    def test_rounded_float_to_int(self, rounding_fn):
+        # The base variable is already supported on the integers, so the cast only
+        # relabels the dtype and must not introduce a second truncation
+        x = pt.random.normal(0.5, 1)
+        y = pt.cast(rounding_fn(x), "int64")
+        y_vv = y.clone()
+
+        np.testing.assert_allclose(
+            logp(y, y_vv).eval({y_vv: 1}),
+            logp(rounding_fn(x), pt.constant(1.0)).eval(),
+        )
+
+    @pytest.mark.parametrize("out_dtype", ["bool", "uint8"])
+    def test_non_truncating_discretizing_cast_not_measurable(self, out_dtype):
+        # Casting to bool tests `x != 0` and casting to an unsigned int wraps around
+        # for negative values; neither is the truncation that float -> int performs
         y = pt.cast(pt.random.normal(), out_dtype)
         with pytest.raises(NotImplementedError):
             logp(y, y.clone())