From 657cb0510f929039d4ab0584ed92302609d5dd34 Mon Sep 17 00:00:00 2001 From: Ricardo Vieira Date: Sun, 9 Aug 2026 22:16:48 +0200 Subject: [PATCH] PatternNodeRewriter: Allow matching variants --- pytensor/graph/rewriting/basic.py | 195 ++++++++++++++++++++-------- pytensor/graph/rewriting/unify.py | 39 +++++- pytensor/tensor/rewriting/math.py | 72 ++++------ tests/graph/rewriting/test_basic.py | 58 ++++++++- tests/tensor/rewriting/test_math.py | 34 ++++- 5 files changed, 289 insertions(+), 109 deletions(-) diff --git a/pytensor/graph/rewriting/basic.py b/pytensor/graph/rewriting/basic.py index 5b31301285..0183dc9437 100644 --- a/pytensor/graph/rewriting/basic.py +++ b/pytensor/graph/rewriting/basic.py @@ -33,6 +33,9 @@ PatternVar, convert_strs_to_vars, match_pattern, + matches_op, + pattern_anchor_depths, + pattern_tokens, reify_pattern, ) from pytensor.graph.traversal import ( @@ -1452,9 +1455,16 @@ class PatternNodeRewriter(NodeRewriter): place. The input pattern cannot just be a string but the output pattern can. - If you put a constant variable in the input pattern, there will be a - match iff a constant variable with the same value and the same type - is found in its place. + ``in_pattern`` may also be a list of alternative input patterns, tried in + order with the first match winning. All alternatives must use the same + variables, which they share with the single output pattern. This is useful + to cover equivalent spellings that different canonicalization stages + produce, e.g. ``1 - x``, ``1 + (-x)`` and ``1 + (-1) * x``. + + If you put a constant variable (or a raw number) in the input pattern, + there will be a match iff a constant variable with the same value is found + in its place. The comparison is broadcast-tolerant and ignores dtype: a + literal ``1`` matches a float64 scalar ``1.0`` as well as a vector of ones. You can add a constraint to the match by using the ``dict(...)`` form described above with a ``'constraint'`` key. The constraint must be a @@ -1543,12 +1553,11 @@ def output_fn(fgraph, node, s): def __init__( self, - in_pattern: tuple, + in_pattern: tuple | list[tuple], out_pattern: tuple | Callable | str, allow_multiple_clients: bool = False, name: str | None = None, tracks=(), - get_nodes=None, values_eq_approx=None, allow_cast=True, ): @@ -1557,7 +1566,9 @@ def __init__( Parameters ---------- in_pattern - The input pattern that we want to replace. + The input pattern that we want to replace, or a list of equivalent + input patterns tried in order (first match wins). All patterns in a + list must use the same variables. out_pattern The replacement pattern. Or a callable that takes (fgraph, node, subs_dict) as inputs, and returns the replacement variable (or None/False to reject the rewrite). @@ -1567,11 +1578,10 @@ def __init__( name Set the name of this rewriter. tracks - The values that :meth:`self.tracks` will return. - get_nodes - If you provide `tracks`, you must provide this parameter. It must be a - function that takes the tracked node and returns a list of nodes on - which we will try this rewrite. + Op instances that trigger this rewrite. Each tracked Op must appear + in the input pattern(s); when it appears below the root, candidate + root nodes are the tracked node's clients at the tracked op's depth + in the pattern. values_eq_approx If specified, this value will be assigned to the ``values_eq_approx`` tag of the output variable. This is used by DebugMode to determine if rewrites are correct. @@ -1580,63 +1590,139 @@ def __init__( Notes ----- - `tracks` and `get_nodes` can be used to make this rewrite track a less - frequent `Op`, which will prevent the rewrite from being tried as often. + `tracks` can be used to make this rewrite track a less frequent `Op` + inside the pattern, which will prevent the rewrite from being tried as + often. """ var_map: dict[str, PatternVar] = {} - self.in_pattern = convert_strs_to_vars(in_pattern, var_map=var_map) + in_patterns = in_pattern if isinstance(in_pattern, list) else [in_pattern] + self.in_patterns = [ + convert_strs_to_vars(p, var_map=var_map) for p in in_patterns + ] + for pattern, orig in zip(self.in_patterns, in_patterns, strict=True): + if not isinstance(pattern, PatternNode): + raise TypeError( + "Each in_pattern must be a tuple starting with an Op or OpPattern; " + f"got {orig!r} of type {type(orig)}" + ) self.out_pattern = convert_strs_to_vars(out_pattern, var_map=var_map) - if not isinstance(self.in_pattern, PatternNode): - raise TypeError( - "The in_pattern must be a tuple starting with an Op or OpPattern; " - f"got {in_pattern!r} of type {type(in_pattern)}" - ) + tokens = pattern_tokens(self.in_patterns[0]) + for pattern, orig in zip(self.in_patterns[1:], in_patterns[1:], strict=True): + if pattern_tokens(pattern) != tokens: + raise ValueError( + "All in_pattern alternatives must use the same variables; " + f"{orig!r} uses {sorted(pattern_tokens(pattern))} " + f"while {in_patterns[0]!r} uses {sorted(tokens)}" + ) self.values_eq_approx = values_eq_approx self.allow_cast = allow_cast self.allow_multiple_clients = allow_multiple_clients if name: self.__name__ = name - self.get_nodes = get_nodes - if tracks != (): - if not get_nodes: - raise ValueError("Custom `tracks` requires `get_nodes` to be provided.") + + tracks = list(tracks) + for track in tracks: + if not isinstance(track, Op): + raise TypeError( + f"PatternNodeRewriter tracks entries must be Op instances; got {track!r}" + ) + # Anchor each pattern at its root and at tracked sub-patterns; transform + # finds candidate roots by walking up an anchored node's clients to the + # anchor's depth in the pattern. + self._anchors = [] + for pattern in self.in_patterns: + self._anchors.append( + [ + (head, depth) + for head, depth in pattern_anchor_depths(pattern) + if depth == 0 or any(matches_op(head, track) for track in tracks) + ] + ) + if tracks: + for track in tracks: + if not any( + matches_op(head, track) + for anchors in self._anchors + for head, _ in anchors + ): + raise ValueError( + f"Tracked op {track} does not appear in any in_pattern" + ) self._tracks = tracks else: - op = self.in_pattern.op_match - if isinstance(op, Op): - self._tracks = [op] - elif isinstance(op, OpPattern): - self._tracks = [op.op_type] - else: - raise ValueError( - f"The in_pattern must start with a specific Op or an OpPattern instance. " - f"Got {op!r}, with type {type(op)}." - ) + roots = [] + for pattern in self.in_patterns: + op = pattern.op_match + entry = op.op_type if isinstance(op, OpPattern) else op + if entry not in roots: + roots.append(entry) + self._tracks = roots def tracks(self): return self._tracks - def transform(self, fgraph, node, enforce_tracks: bool = False, get_nodes=True): - """Check if the graph from node corresponds to ``in_pattern``. + def transform(self, fgraph, node, enforce_tracks: bool = False): + """Check if the graph from node corresponds to one of the ``in_patterns``. If it does, it constructs ``out_pattern`` and performs the replacement. + The match is anchored at ``node``: it must correspond to a pattern's + root or to a tracked sub-pattern. The node's clients are walked once up + to the deepest anchor; each pattern is then matched against its own + candidate roots (the levels at its anchored depths, deduplicated), in + pattern order. """ - if get_nodes and self.get_nodes is not None: - for real_node in self.get_nodes(fgraph, node): - ret = self.transform(fgraph, real_node, get_nodes=False) - if ret is not False and ret is not None: - return dict(zip(real_node.outputs, ret, strict=True)) - - if len(node.outputs) != 1: - # PatternNodeRewriter doesn't support replacing multi-output nodes + node_op = node.op + # For each pattern, the depths of its anchors that match the fired node + applicable = [ + (pattern, {depth for head, depth in anchors if matches_op(head, node_op)}) + for pattern, anchors in zip(self.in_patterns, self._anchors, strict=True) + ] + max_depth = max( + (max(depths) for _, depths in applicable if depths), default=None + ) + if max_depth is None: return False - s = match_pattern(self.in_pattern, node) - if s is None: - return False + clients = fgraph.clients + levels = [(node,)] + for _ in range(max_depth): + level = { + client: None + for level_node in levels[-1] + for out in level_node.outputs + for client, _ in clients[out] + if not isinstance(client.op, Output) + } + if not level: + break + levels.append(level) + + for pattern, depths in applicable: + candidates = { + root: None + for depth in sorted(depths) + if depth < len(levels) + for root in levels[depth] + } + for root in candidates: + if len(root.outputs) != 1: + # PatternNodeRewriter doesn't support replacing multi-output nodes + continue + s = match_pattern(pattern, root) + if s is None: + continue + ret = self._rewrite_from_match(fgraph, root, s) + if ret is None: + continue + if root is node: + return [ret] + return {root.outputs[0]: ret} + return False + def _rewrite_from_match(self, fgraph, node, s): + """Build the replacement variable for a matched root ``node``, or None if rejected.""" if not self.allow_multiple_clients: # The matched subgraph's internal nodes (those strictly between the # captured inputs and the root) must not be used elsewhere, or the @@ -1653,14 +1739,14 @@ def transform(self, fgraph, node, enforce_tracks: bool = False, get_nodes=True): for v in vars_between(input_vars, node.inputs) if v not in input_vars ): - return False + return None if callable(self.out_pattern): # token is the variable name used in the original pattern ret = self.out_pattern(fgraph, node, {k.token: v for k, v in s.items()}) if ret is None or ret is False: # The output function is still allowed to reject the rewrite - return False + return None if not isinstance(ret, Variable): raise ValueError( f"The output of the PatternNodeRewriter callable must be a variable got {ret} of type {type(ret)}." @@ -1681,15 +1767,15 @@ def transform(self, fgraph, node, enforce_tracks: bool = False, get_nodes=True): and isinstance(old_out.type, TensorType) and isinstance(ret.type, TensorType) ): - return False + return None # Try to cast tensors ret = ret.astype(old_out.type.dtype) if not old_out.type.is_super(ret.type): # Still doesn't match - return False + return None - return [ret] + return ret def __str__(self): if getattr(self, "__name__", None): @@ -1708,9 +1794,8 @@ def pattern_to_str(pattern): else: return str(pattern) - return ( - f"{pattern_to_str(self.in_pattern)} -> {pattern_to_str(self.out_pattern)}" - ) + in_str = " | ".join(pattern_to_str(p) for p in self.in_patterns) + return f"{in_str} -> {pattern_to_str(self.out_pattern)}" def __repr__(self): return str(self) @@ -1718,7 +1803,7 @@ def __repr__(self): def print_summary(self, stream=sys.stdout, level=0, depth=-1): name = getattr(self, "__name__", getattr(self, "name", None)) print( - f"{' ' * level}{self.__class__.__name__} {name}({self.in_pattern}, {self.out_pattern}) id={id(self)}", + f"{' ' * level}{self.__class__.__name__} {name}({self.in_patterns}, {self.out_pattern}) id={id(self)}", file=stream, ) diff --git a/pytensor/graph/rewriting/unify.py b/pytensor/graph/rewriting/unify.py index 27a25d4155..99e1320eed 100644 --- a/pytensor/graph/rewriting/unify.py +++ b/pytensor/graph/rewriting/unify.py @@ -1,4 +1,4 @@ -from collections.abc import Callable, Mapping, Sequence +from collections.abc import Callable, Generator, Mapping, Sequence from dataclasses import dataclass from numbers import Number from types import UnionType @@ -342,6 +342,43 @@ def _convert(y, op_prop: bool = False): return _convert(x) +def matches_op(op_match, op: Op) -> bool: + """Whether a pattern head (Op or OpPattern) can match ``op``.""" + if isinstance(op_match, OpPattern): + return op_match.match_op(op) + return bool(op_match == op) + + +def pattern_tokens(pattern) -> frozenset[str]: + """Tokens of all PatternVars and Asterisks in a converted pattern tree.""" + if isinstance(pattern, PatternVar | Asterisk): + return frozenset((pattern.token,)) + if isinstance(pattern, PatternNode): + return pattern_tokens(pattern.op_match).union( + *(pattern_tokens(inp) for inp in pattern.inputs) + ) + if isinstance(pattern, OpPattern): + return frozenset().union( + *(pattern_tokens(param) for _, param in pattern.parameters) + ) + return frozenset() + + +def pattern_anchor_depths( + pattern: PatternNode, depth: int = 0 +) -> Generator[tuple[Any, int], None, None]: + """Enumerate ``(head, depth)`` pairs for every PatternNode in the tree. + + ``head`` is the sub-pattern's Op (or OpPattern) and ``depth`` its distance + from the pattern root. A node matching ``head`` can reach candidate pattern + roots by walking up through its clients ``depth`` times. + """ + yield pattern.op_match, depth + for inp in pattern.inputs: + if isinstance(inp, PatternNode): + yield from pattern_anchor_depths(inp, depth + 1) + + def match_pattern( pattern: PatternNode, node, diff --git a/pytensor/tensor/rewriting/math.py b/pytensor/tensor/rewriting/math.py index bb62d2a1ee..39ed65ac35 100644 --- a/pytensor/tensor/rewriting/math.py +++ b/pytensor/tensor/rewriting/math.py @@ -2,7 +2,7 @@ import itertools from collections import defaultdict -from functools import partial, reduce +from functools import reduce import numpy as np @@ -3003,9 +3003,6 @@ def local_greedy_distributor(fgraph, node): return [rval] -get_clients_at_depth1 = partial(get_clients_at_depth, depth=1) -get_clients_at_depth2 = partial(get_clients_at_depth, depth=2) - # 1+erf(x)=>erfc(-x) local_one_plus_erf = PatternNodeRewriter( (add, 1, (erf, "x")), @@ -3013,7 +3010,6 @@ def local_greedy_distributor(fgraph, node): allow_multiple_clients=True, name="local_one_plus_erf", tracks=[erf], - get_nodes=get_clients_at_depth1, ) register_canonicalize(local_one_plus_erf) register_stabilize(local_one_plus_erf) @@ -3026,7 +3022,6 @@ def local_greedy_distributor(fgraph, node): allow_multiple_clients=True, name="local_one_minus_erf", tracks=[erf], - get_nodes=get_clients_at_depth1, ) register_canonicalize(local_one_minus_erf) register_stabilize(local_one_minus_erf) @@ -3041,7 +3036,6 @@ def local_greedy_distributor(fgraph, node): allow_multiple_clients=True, name="local_erf_minus_one", tracks=[erf], - get_nodes=get_clients_at_depth1, ) register_canonicalize(local_erf_minus_one) register_stabilize(local_erf_minus_one) @@ -3054,7 +3048,6 @@ def local_greedy_distributor(fgraph, node): allow_multiple_clients=True, name="local_one_minus_erfc", tracks=[erfc], - get_nodes=get_clients_at_depth1, ) register_canonicalize(local_one_minus_erfc) register_stabilize(local_one_minus_erfc) @@ -3067,7 +3060,6 @@ def local_greedy_distributor(fgraph, node): allow_multiple_clients=True, name="local_erf_neg_minus_one", tracks=[erfc], - get_nodes=get_clients_at_depth1, ) register_canonicalize(local_erf_neg_minus_one) register_stabilize(local_erf_neg_minus_one) @@ -3320,51 +3312,31 @@ def isclose(x, ref, rtol=0, atol=0, num_ulps=10): return np.allclose(x, ref, rtol=rtol, atol=atol) -def _is_1(expr): - """ - - Returns - ------- - bool - True iff expr is a constant close to 1. - - """ - try: - v = get_underlying_scalar_constant_value(expr) - return isclose(v, 1) - except NotScalarConstantError: - return False - - logsigm_to_softplus = PatternNodeRewriter( (log, (sigmoid, "x")), (neg, (softplus, (neg, "x"))), allow_multiple_clients=True, values_eq_approx=values_eq_approx_remove_inf, tracks=[sigmoid], - get_nodes=get_clients_at_depth1, ) +# log(1 - sigmoid(x)) / log1p(-sigmoid(x)) -> -softplus(x), covering both the +# as-written spellings and the canonicalized ones (neg -> mul(-1), sub -> add) log1msigm_to_softplus = PatternNodeRewriter( - (log, (sub, dict(pattern="y", constraint=_is_1), (sigmoid, "x"))), + [ + (log, (sub, 1, (sigmoid, "x"))), + (log, (add, 1, (mul, -1, (sigmoid, "x")))), + (log1p, (neg, (sigmoid, "x"))), + (log1p, (mul, -1, (sigmoid, "x"))), + ], (neg, (softplus, "x")), allow_multiple_clients=True, values_eq_approx=values_eq_approx_remove_inf, tracks=[sigmoid], - get_nodes=get_clients_at_depth2, -) -log1p_neg_sigmoid = PatternNodeRewriter( - (log1p, (neg, (sigmoid, "x"))), - (neg, (softplus, "x")), - values_eq_approx=values_eq_approx_remove_inf, - allow_multiple_clients=True, - tracks=[sigmoid], - get_nodes=get_clients_at_depth2, ) register_stabilize(logsigm_to_softplus, name="logsigm_to_softplus") register_stabilize(log1msigm_to_softplus, name="log1msigm_to_softplus") -register_stabilize(log1p_neg_sigmoid, name="log1p_neg_sigmoid") -register_specialize(log1p_neg_sigmoid, name="log1p_neg_sigmoid") +register_specialize(log1msigm_to_softplus, name="log1msigm_to_softplus") def is_1pexp(t, only_process_constants=True): @@ -3901,10 +3873,12 @@ def local_reciprocal_1_plus_exp(fgraph, node): # 1 - sigmoid(x) -> sigmoid(-x) local_1msigmoid = PatternNodeRewriter( - (sub, dict(pattern="y", constraint=_is_1), (sigmoid, "x")), + [ + (sub, 1, (sigmoid, "x")), + (add, 1, (mul, -1, (sigmoid, "x"))), + ], (sigmoid, (neg, "x")), tracks=[sigmoid], - get_nodes=get_clients_at_depth1, name="local_1msigmoid", ) register_stabilize(local_1msigmoid) @@ -3974,10 +3948,14 @@ def logmexpm1_to_log1mexp(fgraph, node): # Composed with log(exp(x)) -> x this also covers logit(sigmoid(x)) -> x, so no # separate rewrite is needed for the logged form. local_odds_sigmoid = PatternNodeRewriter( - (true_div, (sigmoid, "x"), (sub, 1, (sigmoid, "x"))), + [ + (true_div, (sigmoid, "x"), (sub, 1, (sigmoid, "x"))), + (true_div, (sigmoid, "x"), (add, 1, (mul, -1, (sigmoid, "x")))), + (true_div, (sigmoid, "x"), (sigmoid, (neg, "x"))), + (true_div, (sigmoid, "x"), (sigmoid, (mul, -1, "x"))), + ], (exp, "x"), tracks=[sigmoid], - get_nodes=get_clients_at_depth1, allow_multiple_clients=True, name="local_odds_sigmoid", ) @@ -3987,10 +3965,14 @@ def logmexpm1_to_log1mexp(fgraph, node): # (1 - sigmoid(x)) / sigmoid(x) -> exp(-x) local_inv_odds_sigmoid = PatternNodeRewriter( - (true_div, (sub, 1, (sigmoid, "x")), (sigmoid, "x")), + [ + (true_div, (sub, 1, (sigmoid, "x")), (sigmoid, "x")), + (true_div, (add, 1, (mul, -1, (sigmoid, "x"))), (sigmoid, "x")), + (true_div, (sigmoid, (neg, "x")), (sigmoid, "x")), + (true_div, (sigmoid, (mul, -1, "x")), (sigmoid, "x")), + ], (exp, (neg, "x")), tracks=[sigmoid], - get_nodes=get_clients_at_depth1, allow_multiple_clients=True, name="local_inv_odds_sigmoid", ) @@ -4036,7 +4018,6 @@ def local_useless_conj(fgraph, node): name="local_log_kv", # Start the rewrite from the less likely kve node tracks=[kve], - get_nodes=get_clients_at_depth2, ) register_stabilize(local_log_kv) @@ -4050,7 +4031,6 @@ def local_useless_conj(fgraph, node): name="local_log_iv", # Start the rewrite from the less likely ive node tracks=[ive], - get_nodes=get_clients_at_depth2, ) register_stabilize(local_log_iv) diff --git a/tests/graph/rewriting/test_basic.py b/tests/graph/rewriting/test_basic.py index d69249b41c..c40b84828b 100644 --- a/tests/graph/rewriting/test_basic.py +++ b/tests/graph/rewriting/test_basic.py @@ -322,6 +322,61 @@ def test_op_pattern(self): rewriter.rewrite(g) assert equal_computations(g.outputs, [e1, a]) + def test_multiple_in_patterns(self): + rewriter = PatternNodeRewriter( + [(op1, "1"), (op2, (op3, "1"))], + (op4, "1"), + ) + assert rewriter.tracks() == [op1, op2] + x = MyVariable("x") + for e in (op1(x), op2(op3(x))): + g = FunctionGraph([x], [e]) + EquilibriumGraphRewriter([rewriter], max_use_ratio=1).rewrite(g) + assert str(g) == "FunctionGraph(Op4(x))" + + def test_tracked_subpattern_anchoring(self): + # tracks below the pattern root: candidate roots are found by walking + # up from the tracked node's clients along the pattern structure + rewriter = PatternNodeRewriter( + [ + (op1, (op5, "1")), + (op2, (op3, (op5, "1"))), + ], + (op4, "1"), + tracks=[op5], + ) + assert rewriter.tracks() == [op5] + x = MyVariable("x") + for e in (op1(op5(x)), op2(op3(op5(x)))): + g = FunctionGraph([x], [e]) + EquilibriumGraphRewriter([rewriter], max_use_ratio=1).rewrite(g) + assert str(g) == "FunctionGraph(Op4(x))" + + # A direct call on the pattern root node also works + e = op1(op5(x)) + g = FunctionGraph([x], [e], clone=False) + assert rewriter.transform(g, e.owner) is not False + + # A call on the tracked node returns a dict keyed by the root's output + e = op2(op3(op5(x))) + g = FunctionGraph([x], [e], clone=False) + op5_node = e.owner.inputs[0].owner.inputs[0].owner + repl = rewriter.transform(g, op5_node) + assert isinstance(repl, dict) + [(replaced_out, new_out)] = repl.items() + assert replaced_out is g.outputs[0] + assert equal_computations([new_out], [op4(x)]) + + def test_in_patterns_validation(self): + with pytest.raises(ValueError, match="same variables"): + PatternNodeRewriter([(op1, "1"), (op2, "2")], (op4, "1")) + + with pytest.raises(ValueError, match="does not appear"): + PatternNodeRewriter((op1, "1"), (op4, "1"), tracks=[op2]) + + with pytest.raises(TypeError, match="Op instances"): + PatternNodeRewriter((op1, "1"), (op4, "1"), tracks=[MyOp]) + class NoInputOp(Op): __props__ = ("param",) @@ -666,7 +721,7 @@ def test_pre_greedy_node_rewriter(): @pytest.mark.parametrize("tracks", [True, False]) @pytest.mark.parametrize("out_pattern", [(op2, "x"), "x", 1.0]) def test_patternsub_values_eq_approx(out_pattern, tracks): - # PatternNodeRewriter would fail when `values_eq_approx` and `get_nodes` were specified + # PatternNodeRewriter would fail when `values_eq_approx` and `tracks` were specified x = MyVariable("x") e = op1(x) fg = FunctionGraph([x], [e], clone=False) @@ -677,7 +732,6 @@ def test_patternsub_values_eq_approx(out_pattern, tracks): (op1, "x"), out_pattern, tracks=[op1] if tracks else (), - get_nodes=(lambda fgraph, node: [node]) if tracks else None, values_eq_approx=values_eq_approx_always_true, ) ], diff --git a/tests/tensor/rewriting/test_math.py b/tests/tensor/rewriting/test_math.py index 0a2b845052..e0d0654940 100644 --- a/tests/tensor/rewriting/test_math.py +++ b/tests/tensor/rewriting/test_math.py @@ -4358,8 +4358,13 @@ def test_local_1msigmoid(self): (np.array(1.0, "float64") - pt.sigmoid(x), cast(sigmoid(-x), "float64")), (np.array(1.0, "float32") - sigmoid(xd), sigmoid(-xd)), (np.array(1.0, "float64") - sigmoid(xd), sigmoid(-xd)), + # This sum folds to exactly 1.0 in float32 (np.sum(1 / np.array([2, 3, 6], "float32")) - sigmoid(x), sigmoid(-x)), - (np.sum(1 / np.array([2, 3, 6], "float64")) - sigmoid(xd), sigmoid(-xd)), + # In float64 it is one ulp below 1, so it must not be treated as 1 + ( + np.sum(1 / np.array([2, 3, 6], "float64")) - sigmoid(xd), + np.sum(1 / np.array([2, 3, 6], "float64")) - sigmoid(xd), + ), (np.float32(1 - 9e-6) - sigmoid(x), np.float32(1 - 9e-6) - sigmoid(x)), (np.float64(1 - 1e-9) - sigmoid(xd), np.float64(1 - 1e-9) - sigmoid(xd)), ]: @@ -4513,14 +4518,15 @@ def test_log1msigm_to_softplus(self): # assert check_stack_trace(f, ops_to_check='all') f(np.random.random((54, 11)).astype(config.floatX)) - # Test close to 1 + # A constant one ulp away from 1 must not be treated as 1 x_dtype = np.dtype(x.dtype).type out = log(np.nextafter(x_dtype(1), x_dtype(2)) - sigmoid(x)) f = pytensor.function([x], out, mode=self.m) topo = f.maker.fgraph.toposort() - assert len(topo) == 2 - assert isinstance(topo[0].op.scalar_op, pytensor.scalar.Softplus) - assert isinstance(topo[1].op.scalar_op, pytensor.scalar.Neg) + assert not any( + isinstance(getattr(node.op, "scalar_op", None), pytensor.scalar.Softplus) + for node in topo + ) # Same test with a flatten out = log(1 - pt.flatten(sigmoid(x))) @@ -4655,6 +4661,24 @@ def test_local_odds_sigmoid(): result.assert_graph(exp(-1.0 * x)) result.assert_eval(np.array(0.5)) + # the sigmoid(-x) spelling produced by local_1msigmoid + result = RewriteTester([x], [sigmoid(x) / sigmoid(-x)]) + result.assert_graph(exp(x)) + result.assert_eval(np.array(0.5)) + + +def test_1msigmoid_canonicalized_spellings(): + """log1p(-sigmoid(x)) is canonicalized to log1p((-1) * sigmoid(x)) before + stabilize runs, which used to escape log1msigm_to_softplus until specialize. + """ + x = dscalar("x") + + result = RewriteTester( + [x], [log1p(-sigmoid(x))], include=("canonicalize", "stabilize") + ) + result.assert_graph(-softplus(x)) + result.assert_eval(np.array(0.5)) + def test_local_useless_conj(): default_mode = get_default_mode()