Infer logprob of trunc and round_half_away_from_zero rounding

Truncation towards zero is the interval [x, x+1) for positive values,
(x-1, x] for negative ones, and (-1, 1) pooled at zero. Rounding half
away from zero shares the intervals of round half to even, since the
tie-breaking rule only differs on a measure-zero set of the continuous
base variable.

Author: ricardoV94

pymc/logprob/censoring.py

@@ -41,10 +41,27 @@
 from pytensor.graph.basic import Apply
 from pytensor.graph.fg import FunctionGraph
 from pytensor.graph.rewriting.basic import node_rewriter
-from pytensor.scalar.basic import Ceil, Clip, Floor, Maximum, RoundHalfToEven
+from pytensor.scalar.basic import (
+    Ceil,
+    Clip,
+    Floor,
+    Maximum,
+    RoundHalfAwayFromZero,
+    RoundHalfToEven,
+    Trunc,
+)
 from pytensor.scalar.basic import clip as scalar_clip
 from pytensor.tensor import TensorVariable
-from pytensor.tensor.math import ceil, clip, floor, maximum, minimum, round_half_to_even
+from pytensor.tensor.math import (
+    ceil,
+    clip,
+    floor,
+    maximum,
+    minimum,
+    round_half_away_from_zero,
+    round_half_to_even,
+    trunc,
+)
 from pytensor.tensor.variable import TensorConstant
 
 from pymc.logprob.abstract import (
@@ -267,10 +284,10 @@ def clip_icdf(op, value, base_rv, lower_bound, upper_bound, **kwargs):
 class MeasurableRound(MeasurableElemwise):
     """A placeholder used to specify a log-likelihood for a clipped RV sub-graph."""
 
-    valid_scalar_types = (RoundHalfToEven, Floor, Ceil)
+    valid_scalar_types = (RoundHalfToEven, RoundHalfAwayFromZero, Floor, Ceil, Trunc)
 
 
-@node_rewriter(tracks=[ceil, floor, round_half_to_even])
+@node_rewriter(tracks=[ceil, floor, round_half_to_even, round_half_away_from_zero, trunc])
 def find_measurable_roundings(fgraph: FunctionGraph, node: Apply) -> list[TensorVariable] | None:
     if not filter_measurable_variables(node.inputs):
         return None
@@ -315,10 +332,16 @@ def round_logprob(op, values, base_rv, **kwargs):
             0 & \text{otherwise},
         \end{cases}
 
+    The probability of a distribution rounded towards zero is given by the
+    rounded-down probability for positive values and the rounded-up probability
+    for negative values, with both intervals pooled at zero.
+
     """
     (value,) = values
 
-    if isinstance(op.scalar_op, RoundHalfToEven):
+    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
         value = pt.round(value)
         value_upper = value + 0.5
         value_lower = value - 0.5
@@ -330,6 +353,13 @@ def round_logprob(op, values, base_rv, **kwargs):
         value = pt.ceil(value)
         value_upper = value
         value_lower = value - 1.0
+    elif isinstance(op.scalar_op, Trunc):
+        # Truncation rounds towards zero: [x, x+1) for x >= 0, (x-1, x] for x < 0,
+        # and (-1, 1) for x == 0 (open/closed bounds are equivalent for the
+        # continuous base variables this rewrite applies to)
+        value = pt.trunc(value)
+        value_upper = value + (value >= 0)
+        value_lower = value - (value <= 0)
     else:
         raise TypeError(f"Unsupported scalar_op {op.scalar_op}")  # pragma: no cover
 

tests/logprob/test_censoring.py

@@ -235,7 +235,9 @@ def test_clip_transform():
     assert np.isclose(obs_logp, exp_logp)
 
 
-@pytest.mark.parametrize("rounding_op", (pt.round, pt.floor, pt.ceil))
+@pytest.mark.parametrize(
+    "rounding_op", (pt.round, pt.round_half_away_from_zero, pt.floor, pt.ceil, pt.trunc)
+)
 def test_rounding(rounding_op):
     loc = 1
     scale = 2
@@ -250,12 +252,16 @@ def test_rounding(rounding_op):
     assert logprob is not None
 
     x_sp = st.norm(loc, scale)
-    if rounding_op == pt.round:
+    if rounding_op in (pt.round, pt.round_half_away_from_zero):
         expected_logp = np.log(x_sp.cdf(test_value + 0.5) - x_sp.cdf(test_value - 0.5))
     elif rounding_op == pt.floor:
         expected_logp = np.log(x_sp.cdf(test_value + 1.0) - x_sp.cdf(test_value))
     elif rounding_op == pt.ceil:
         expected_logp = np.log(x_sp.cdf(test_value) - x_sp.cdf(test_value - 1.0))
+    elif rounding_op == pt.trunc:
+        expected_logp = np.log(
+            x_sp.cdf(test_value + (test_value >= 0)) - x_sp.cdf(test_value - (test_value <= 0))
+        )
     else:
         raise NotImplementedError()