diff --git a/pytensor/scalar/basic.py b/pytensor/scalar/basic.py index e3c088ff81..a8f6a8ab64 100644 --- a/pytensor/scalar/basic.py +++ b/pytensor/scalar/basic.py @@ -1296,6 +1296,11 @@ class UnaryScalarOp(ScalarOp): monotonic_increasing = False monotonic_decreasing = False + # Set only when the output is >= 0 for every real input, including the sign of zero + # (`sqrt` returns -0.0 for -0.0, so it does not qualify). Consumers must still check + # the dtype: signed integers wrap on overflow and complex values are unordered. + non_negative = False + def c_code_contiguous(self, node, name, inputs, outputs, sub): (x,) = inputs (z,) = outputs @@ -2523,6 +2528,7 @@ def cast(x, dtype): class Abs(UnaryScalarOp): preserves_zero = True + non_negative = True nfunc_spec = ("abs", 1, 1) def make_node(self, x): @@ -3084,6 +3090,7 @@ def c_code(self, node, name, inputs, outputs, sub): class Exp(UnaryScalarOp): monotonic_increasing = True + non_negative = True nfunc_spec = ("exp", 1, 1) amd_float32 = "amd_vrsa_expf" amd_float64 = "amd_vrda_exp" @@ -3123,6 +3130,7 @@ def c_code(self, node, name, inputs, outputs, sub): class Exp2(UnaryScalarOp): monotonic_increasing = True + non_negative = True nfunc_spec = ("exp2", 1, 1) def impl(self, x): @@ -3201,6 +3209,7 @@ def c_code_cache_version(self): class Sqr(UnaryScalarOp): preserves_zero = True + non_negative = True nfunc_spec = ("square", 1, 1) def impl(self, x): @@ -3546,6 +3555,8 @@ class Cosh(UnaryScalarOp): """ + non_negative = True + nfunc_spec = ("cosh", 1, 1) def impl(self, x): diff --git a/pytensor/scalar/math.py b/pytensor/scalar/math.py index f13297e572..d36f91054e 100644 --- a/pytensor/scalar/math.py +++ b/pytensor/scalar/math.py @@ -90,6 +90,7 @@ def c_code(self, node, name, inp, out, sub): class Erfc(UnaryScalarOp): monotonic_decreasing = True + non_negative = True nfunc_spec = ("scipy.special.erfc", 1, 1) def impl(self, x): @@ -1170,6 +1171,7 @@ class Sigmoid(UnaryScalarOp): """ monotonic_increasing = True + non_negative = True nfunc_spec = ("scipy.special.expit", 1, 1) def impl(self, x): @@ -1225,6 +1227,7 @@ class Softplus(UnaryScalarOp): """ monotonic_increasing = True + non_negative = True def impl(self, x): # If x is an int8 or uint8, numpy.exp will compute the result in diff --git a/pytensor/tensor/rewriting/linalg/summary.py b/pytensor/tensor/rewriting/linalg/summary.py index 2e41fb3d1b..ae8a4e422c 100644 --- a/pytensor/tensor/rewriting/linalg/summary.py +++ b/pytensor/tensor/rewriting/linalg/summary.py @@ -10,7 +10,7 @@ copy_stack_trace, node_rewriter, ) -from pytensor.scalar.basic import Abs, Exp, Log, Sign, Sqr +from pytensor.scalar.basic import Abs, Log, Sign, Sqr from pytensor.tensor.basic import ones from pytensor.tensor.blockwise import Blockwise from pytensor.tensor.elemwise import Elemwise @@ -27,13 +27,14 @@ register_stabilize, ) from pytensor.tensor.rewriting.linalg.utils import matrix_diagonal_product +from pytensor.tensor.subtensor import _is_provably_non_negative @register_stabilize @register_specialize @node_rewriter([log]) def local_log_prod_to_sum_log(fgraph, node): - """Rewrite log(prod(x)) as sum(log(x)), when x is known to be positive.""" + """Rewrite log(prod(x)) as sum(log(x)), when x is known to be non-negative.""" [p] = node.inputs match p.owner_op_and_inputs: case (Prod(axis=axis), x): @@ -41,11 +42,7 @@ def local_log_prod_to_sum_log(fgraph, node): # returns the sign of the prod multiplication. # TODO: The product of diagonals of a Cholesky(A) are also strictly positive - match x.owner_op: - case Elemwise(Abs() | Sqr() | Exp()): - return [log(x).sum(axis=axis)] - - if getattr(x.tag, "positive", False): + if _is_provably_non_negative(x): return [log(x).sum(axis=axis)] # Special case for log(abs(prod(x))) -> sum(log(abs(x))) that shows up in slogdet diff --git a/pytensor/tensor/rewriting/math.py b/pytensor/tensor/rewriting/math.py index bb62d2a1ee..6005412b1a 100644 --- a/pytensor/tensor/rewriting/math.py +++ b/pytensor/tensor/rewriting/math.py @@ -110,7 +110,11 @@ from pytensor.tensor.rewriting.blockwise import blockwise_of from pytensor.tensor.rewriting.elemwise import apply_local_dimshuffle_lift from pytensor.tensor.shape import Shape, Shape_i, specify_shape -from pytensor.tensor.subtensor import Subtensor, _is_provably_positive +from pytensor.tensor.subtensor import ( + Subtensor, + _is_provably_non_negative, + _is_provably_positive, +) from pytensor.tensor.type import ( complex_dtypes, uint_dtypes, @@ -594,22 +598,43 @@ def local_sqrt_sqr(fgraph, node): return [new_out] +@register_canonicalize @register_specialize -@node_rewriter([log]) -def local_log_sqrt(fgraph, node): - x = node.inputs[0] +@node_rewriter([pt_abs]) +def local_useless_abs(fgraph, node): + # Case for abs(x) -> x, when x is already non-negative + [x] = node.inputs - if ( - not x.owner - or not isinstance(x.owner.op, Elemwise) - or not isinstance(x.owner.op.scalar_op, ps.Sqrt) - ): + if not _is_provably_non_negative(x): return - # Case for log(sqrt(x)) -> 0.5 * log(x) - x = x.owner.inputs[0] + return [x] + + +@register_stabilize +@register_specialize +@node_rewriter([log]) +def local_log_sqrt_sqr(fgraph, node): + [x] = node.inputs + + match x.owner_op_and_inputs: + # Case for log(sqrt(x)) -> 0.5 * log(x) + case (Elemwise(ps.Sqrt()), inner): + factor, inner_out = 0.5, log(inner) + + # Case for log(sqr(x)) -> 2 * log(abs(x)), which never materializes the square, + # so a value that would over- or underflow when squared survives the log. + case (Elemwise(ps.Sqr()), inner): + # abs() of a complex input is real and would drop the imaginary part + if inner.dtype.startswith("complex"): + return + factor, inner_out = 2.0, log(pt_abs(inner)) + + case _: + return + old_out = node.outputs[0] - new_out = mul(as_tensor_variable(0.5, dtype=x.dtype), log(x)) + new_out = mul(as_tensor_variable(factor, dtype=old_out.dtype), inner_out) if new_out.dtype != old_out.dtype: new_out = cast(new_out, old_out.dtype) diff --git a/pytensor/tensor/subtensor.py b/pytensor/tensor/subtensor.py index 62a9879399..31cb18d726 100644 --- a/pytensor/tensor/subtensor.py +++ b/pytensor/tensor/subtensor.py @@ -22,6 +22,7 @@ ScalarMaximum, ScalarMinimum, ScalarVariable, + UnaryScalarOp, ) from pytensor.tensor import ( TensorLike, @@ -253,12 +254,15 @@ def _is_provably_positive(var, strict: bool = True) -> bool: - ``minimum(a, b)`` when both ``a`` and ``b`` are positive. - ``maximum(a, b)`` when at least one of ``a``, ``b`` is positive. - Three further cases prove non-negativity but not strict positivity, so they + Four further cases prove non-negativity but not strict positivity, so they are recognized only when ``strict=False``: - Unsigned-integer dtype (a ``uint`` may be 0). - ``Shape`` / ``Shape_i`` outputs (a dimension may be 0). - ``Cast`` of a non-negative input (a float :math:`0 < x < 1` truncates to 0). + - A unary scalar op declaring ``non_negative`` (``abs``, ``sqr``, ``exp``, ...), + restricted to float and unsigned outputs: signed integers wrap on overflow + (``sqr(int8(12)) == -112``) and complex values are unordered. Parameters ---------- @@ -293,6 +297,13 @@ def _is_provably_positive(var, strict: bool = True) -> bool: scalar_op = op.scalar_op if not strict and isinstance(scalar_op, Cast): return _is_provably_positive(var.owner.inputs[0], strict) + if ( + not strict + and isinstance(scalar_op, UnaryScalarOp) + and scalar_op.non_negative + and var.type.dtype.startswith(("float", "uint")) + ): + return True if isinstance(scalar_op, ScalarMinimum): return all(_is_provably_positive(i, strict) for i in var.owner.inputs) if isinstance(scalar_op, ScalarMaximum): diff --git a/tests/tensor/rewriting/linalg/test_summary.py b/tests/tensor/rewriting/linalg/test_summary.py index 517e178ebc..f54c9f5926 100644 --- a/tests/tensor/rewriting/linalg/test_summary.py +++ b/tests/tensor/rewriting/linalg/test_summary.py @@ -10,8 +10,9 @@ from pytensor.tensor.linalg.decomposition import lu, qr, svd from pytensor.tensor.linalg.decomposition.cholesky import cholesky from pytensor.tensor.linalg.summary import Det, SLogDet, det +from pytensor.tensor.math import Prod from pytensor.tensor.type import matrix -from tests.unittest_tools import assert_equal_computations +from tests.unittest_tools import RewriteTester, assert_equal_computations def test_det_of_cholesky(): @@ -239,7 +240,7 @@ def test_slogdet_specialization(): ), pytest.param( lambda x: pt.log(pt.prod(x**2)), - lambda x: pt.sum(pt.log(pt.sqr(x))), + lambda x: np.float64(2.0) * pt.sum(pt.log(pt.abs(x))), id="log_prod_sqr", ), pytest.param( @@ -247,6 +248,11 @@ def test_slogdet_specialization(): lambda x: pt.sum(pt.log(pt.abs(x))), id="log_abs_prod", ), + pytest.param( + lambda x: pt.log(pt.abs(pt.sqr(pt.prod(x)))), + lambda x: np.float64(2.0) * pt.sum(pt.log(pt.abs(x))), + id="log_abs_sqr_prod", + ), pytest.param( lambda x: pt.log(pt.prod(pt.abs(x), axis=0)), lambda x: pt.sum(pt.log(pt.abs(x)), axis=0), @@ -267,30 +273,31 @@ def test_local_log_prod_to_sum_log(original_fn, expected_fn): assert_equal_computations([rewritten], [expected]) -@pytest.mark.parametrize( - "expected, pos_tag", - [ - pytest.param( - lambda x: pt.sum(pt.log(x)), - True, - id="local_log_prod_to_sum_log_positive_tag", - ), - pytest.param( - lambda x: pt.log(pt.prod(x)), - False, - id="local_log_prod_to_sum_log_no_rewrite", - ), - ], -) -def test_local_log_prod_to_sum_log_positive_tag(expected, pos_tag): - x = pt.tensor("x", shape=(3, 4)) - if pos_tag: - x.tag.positive = True +@pytest.mark.parametrize("n", [200, 250, 300]) +def test_log_abs_sqr_prod_no_underflow(n): + """A float32 product of this many terms is zero once squared, so it must not be materialized.""" + x = pt.vector("x", dtype="float32") + out = pt.log(pt.abs(pt.sqr(pt.prod(x)))) + + fn = function([x], out, mode="FAST_RUN") + assert not any(isinstance(node.op, Prod) for node in fn.maker.fgraph.apply_nodes) + + rng = np.random.default_rng(sum(map(ord, "log_sqr_prod"))) + # Mixed signs, so a rewrite that dropped the abs would return nan + x_test = (rng.uniform(0.5, 0.9, size=n) * rng.choice([-1, 1], size=n)).astype( + "float32" + ) + expected = 2 * np.sum(np.log(np.abs(x_test.astype("float64")))) + assert_allclose(fn(x_test), expected, rtol=1e-4) + +def test_local_log_prod_to_sum_log_unknown_sign(): + """A sign-indefinite operand leaves log(prod(x)) alone.""" + x = pt.tensor("x", shape=(3, 4)) out = pt.log(pt.prod(x)) + result = RewriteTester([x], [out], include=["stabilize", "specialize"]) - rewritten = rewrite_graph(out, include=["stabilize", "specialize"]) - assert_equal_computations([rewritten], [expected(x)]) + result.assert_graph(out) @pytest.mark.parametrize( diff --git a/tests/tensor/rewriting/test_math.py b/tests/tensor/rewriting/test_math.py index 0a2b845052..3e4767eb26 100644 --- a/tests/tensor/rewriting/test_math.py +++ b/tests/tensor/rewriting/test_math.py @@ -87,6 +87,7 @@ reciprocal, sigmoid, sign, + sin, sinh, softplus, sqr, @@ -111,12 +112,14 @@ local_div_switch_sink, local_grad_log_erfc_neg, local_greedy_distributor, + local_log_sqrt_sqr, local_mul_canonizer, local_mul_switch_sink, local_neg_to_mul, local_reduce_chain, local_reduce_join, local_sum_prod_of_mul_or_div, + local_useless_abs, mul_canonizer, parse_mul_tree, perform_sigm_times_exp, @@ -2141,17 +2144,145 @@ def test_exp_log_nested(self, nested_expression, expected_switches): assert len(ops_graph) == expected_switches -def test_log_sqrt() -> None: +def test_log_sqrt(): x = pt.tensor("x", shape=(None, None)) - out = log(sqrt(x)) + result = RewriteTester( + [x], [log(sqrt(x))], include=None, custom_rewrite=local_log_sqrt_sqr + ) + + result.assert_graph(0.5 * log(x)) + result.assert_eval(np.array([[1.0, 2.0], [3.0, 4.0]])) + + +def test_log_sqrt_integer_input(): + # A 0.5 factor typed from the integer input would truncate to 0 + x = ivector("x") + result = RewriteTester( + [x], [log(sqrt(x))], include=None, custom_rewrite=local_log_sqrt_sqr + ) + + result.assert_graph(0.5 * log(x)) + result.assert_eval(np.array([2, 3, 4], dtype="int32")) + + +def test_log_sqr(): + x = pt.tensor("x", shape=(None, None)) + result = RewriteTester( + [x], [log(sqr(x))], include=None, custom_rewrite=local_log_sqrt_sqr + ) + + result.assert_graph(2.0 * log(pt_abs(x))) + result.assert_eval(np.array([[1.0, -2.0], [3.0, -4.0]])) + + +@pytest.mark.parametrize( + "inner_fn", [sqr, pt_abs, exp, sigmoid], ids=["sqr", "abs", "exp", "sigmoid"] +) +def test_useless_abs(inner_fn): + x = pt.tensor("x", shape=(None, None)) + result = RewriteTester( + [x], [pt_abs(inner_fn(x))], include=None, custom_rewrite=local_useless_abs + ) + + result.assert_graph(inner_fn(x)) + result.assert_eval(np.array([[1.0, -2.0], [3.0, -4.0]])) - out = rewrite_graph(out, include=["specialize"]) - assert utt.assert_equal_computations( - [out], - [mul(pt.as_tensor_variable([[0.5]], dtype=x.dtype), log(x))], +def test_useless_abs_unsigned_dtype(): + x = pt.vector("x", dtype="uint8") + result = RewriteTester( + [x], [pt_abs(x)], include=None, custom_rewrite=local_useless_abs ) + result.assert_graph(x) + result.assert_eval(np.array([200, 3], dtype="uint8")) + + +def test_useless_abs_of_clipped_input(): + # maximum() is proven non-negative by the shared predicate, not by an op flag + x = pt.vector("x") + result = RewriteTester( + [x], [pt_abs(pt.maximum(x, 0))], include=None, custom_rewrite=local_useless_abs + ) + + result.assert_graph(pt.maximum(x, 0)) + result.assert_eval(np.array([-2.0, 0.0, 3.0])) + + +@pytest.mark.parametrize("inner_fn", [sin, neg], ids=["sin", "neg"]) +def test_useless_abs_sign_indefinite(inner_fn): + x = pt.tensor("x", shape=(None, None)) + out = pt_abs(inner_fn(x)) + result = RewriteTester([x], [out], include=None, custom_rewrite=local_useless_abs) + + result.assert_graph(out) + result.assert_eval(np.array([[1.0, -2.0], [3.0, -4.0]])) + + +def test_useless_abs_keeps_signed_zero(): + # sqrt(-0.0) is -0.0, which abs() normalizes to +0.0, so dropping the abs would flip + # the sign of any downstream division + x = pt.vector("x") + out = pt_abs(sqrt(x)) + result = RewriteTester([x], [out], include=None, custom_rewrite=local_useless_abs) + + result.assert_graph(out) + [rewr_out] = result.rewr_fn(np.array([-0.0])) + assert not np.signbit(rewr_out) + + +def test_useless_abs_signed_integer_overflow(): + # sqr() wraps on signed ints (sqr(int8(12)) == -112), so its output is not + # non-negative and the abs() has to stay + x = pt.vector("x", dtype="int8") + out = pt_abs(sqr(x)) + result = RewriteTester([x], [out], include=None, custom_rewrite=local_useless_abs) + + result.assert_graph(out) + result.assert_eval(np.array([12, 16, 3], dtype="int8")) + + +def test_log_sqr_integer_input(): + x = ivector("x") + result = RewriteTester( + [x], [log(sqr(x))], include=None, custom_rewrite=local_log_sqrt_sqr + ) + + result.assert_graph(2.0 * log(pt_abs(x))) + result.assert_eval(np.array([2, 3, 4], dtype="int32")) + + +def test_log_sqr_extreme_magnitudes(): + # sqr() saturates to inf above ~1e154 and to zero below ~1e-162 in float64, so the + # rewritten graph is deliberately not equivalent to the original one here + x = pt.vector("x") + result = RewriteTester( + [x], [log(sqr(x))], include=None, custom_rewrite=local_log_sqrt_sqr + ) + + x_test = np.array([1e200, 1e-200, 3.0]) + [orig_out] = result.orig_fn(x_test) + [rewr_out] = result.rewr_fn(x_test) + assert np.isinf(orig_out).sum() == 2 + np.testing.assert_allclose(rewr_out, 2 * np.log(np.abs(x_test))) + + +@pytest.mark.parametrize( + "original_fn, rewrite", + [ + pytest.param(lambda x: log(sqr(x)), local_log_sqrt_sqr, id="log_sqr"), + pytest.param(lambda x: pt_abs(sqr(x)), local_useless_abs, id="abs_sqr"), + ], +) +def test_sqr_rewrites_skip_complex(original_fn, rewrite): + # abs() of a complex input is real, so both rewrites would change the output dtype + x = pt.vector("x", dtype="complex128") + out = original_fn(x) + result = RewriteTester([x], [out], include=None, custom_rewrite=rewrite) + + result.assert_graph(out) + result.assert_eval(np.array([1 + 2j, 0.5 - 1j])) + class TestSqrSqrt: def setup_method(self):