diff --git a/pytensor/graph/utils.py b/pytensor/graph/utils.py index d24003230b..67096b7c55 100644 --- a/pytensor/graph/utils.py +++ b/pytensor/graph/utils.py @@ -223,9 +223,11 @@ def __hash__(self): if "__eq__" not in dct: def __eq__(self, other): - return type(self) is type(other) and tuple( - getattr(self, a) for a in props - ) == tuple(getattr(other, a) for a in props) + return self is other or ( + type(self) is type(other) + and tuple(getattr(self, a) for a in props) + == tuple(getattr(other, a) for a in props) + ) dct["__eq__"] = __eq__ diff --git a/pytensor/tensor/elemwise.py b/pytensor/tensor/elemwise.py index 2a3de0a167..05d61aff8d 100644 --- a/pytensor/tensor/elemwise.py +++ b/pytensor/tensor/elemwise.py @@ -382,6 +382,19 @@ def __init__( self.__setstate__(self.__dict__) super().__init__(openmp=openmp) + def __eq__(self, other): + # Hot in rewriting. Identity accepts the singleton-op common case, and + # the scalar_op type check rejects mismatches without paying for the + # props tuples the generated __eq__ would build. + if self is other: + return True + return ( + type(self) is type(other) + and type(self.scalar_op) is type(other.scalar_op) + and self.scalar_op == other.scalar_op + and self.inplace_pattern == other.inplace_pattern + ) + def __getstate__(self): d = copy(self.__dict__) d.pop("ufunc") diff --git a/pytensor/tensor/rewriting/math.py b/pytensor/tensor/rewriting/math.py index 0a7d0dac7e..e2693394b2 100644 --- a/pytensor/tensor/rewriting/math.py +++ b/pytensor/tensor/rewriting/math.py @@ -1,7 +1,6 @@ r"""Rewrites for the `Op`\s in :mod:`pytensor.tensor.math`.""" import itertools -import operator from collections import defaultdict from functools import partial, reduce @@ -36,6 +35,7 @@ moveaxis, ones_like, register_infer_shape, + second, split, switch, zeros, @@ -516,7 +516,12 @@ def local_dot_to_mul(fgraph, node): allow_cast=True, name=f"useless_{op}_of_{inv_op}", ) - register_canonicalize(rewrite) + if op is not neg: + # `local_neg_to_mul` already cancels neg(neg(x)) during + # canonicalize, so matching the pattern there as well is wasted + # work in the equilibrium loop. It is still needed in specialize, + # where `local_neg_to_mul` does not run. + register_canonicalize(rewrite) register_specialize(rewrite) if op is inv_op: @@ -1147,6 +1152,14 @@ class AlgebraicCanonizer(NodeRewriter): should be returned as a list of one element, unless the value is such that ``value = main()``. In that case, the return value should be an empty list. + absorbing_element + The absorbing (annihilating) element of `main`, if it has one, such + that ``main(absorbing_element, x) == absorbing_element`` for every + ``x`` (e.g. ``0`` for mul). When given, any `main` node with a + constant `absorbing_element` among its factors folds to that element + directly, and constants folding to it in `simplify_constants` drop + the surviving non-constant terms. `add` has no absorbing element, so + `local_add_canonizer` leaves this as ``None``. Examples -------- @@ -1172,12 +1185,21 @@ class AlgebraicCanonizer(NodeRewriter): """ - def __init__(self, main, inverse_fn, reciprocal_fn, calculate, use_reciprocal=True): + def __init__( + self, + main, + inverse_fn, + reciprocal_fn, + calculate, + use_reciprocal=True, + absorbing_element=None, + ): self.main = main self.inverse = inverse_fn self.reciprocal = reciprocal_fn self.calculate = calculate self.use_reciprocal = use_reciprocal + self.absorbing_element = absorbing_element self.external_simplifiers = [] @@ -1209,86 +1231,43 @@ def get_num_denum(self, inp): | x * y * z -> ([x, y, z], []) """ - # This function is recursive. The idea is that there is a - # get_num_denum recursion in which the internal ops are all - # one of (main, inverse, reciprocal, DimShuffle) and the - # internal data nodes all have the dtype of the 'input' - # argument. The leaf-Variables of the graph covered by the - # recursion may be of any Variable type. - - if inp.owner is None or inp.owner.op not in [ - self.main, - self.inverse, - self.reciprocal, - ]: - if inp.owner and isinstance(inp.owner.op, DimShuffle): - # If input is a DimShuffle of some input which does - # something like this: - - # * change a vector of length N into a 1xN row matrix - # * change a scalar into a 1x1x1 tensor - # * in general, complete the shape of a tensor - # with broadcastable 1s to the *left* - # Then we will simply discard the DimShuffle and return - # the num/denum of its input - dsn = inp.owner # dimshuffle node - dsop = dsn.op # dimshuffle op - - # the first input of the dimshuffle i.e. the ndarray to redim - dsi0 = dsn.inputs[0] - - # The compatible order is a DimShuffle "new_order" of the form: - # ('x', ..., 'x', 0, 1, 2, ..., dimshuffle_input.type.ndim) - - # That kind of DimShuffle only adds broadcastable - # dimensions on the left, without discarding any - # existing broadcastable dimension and is inserted - # automatically by Elemwise when the inputs have - # different numbers of dimensions (hence why we can - # discard its information - we know we can retrieve it - # later on). - compatible_order = ("x",) * (inp.type.ndim - dsi0.type.ndim) + tuple( - range(dsi0.type.ndim) - ) - if dsop.new_order == compatible_order: - # If the "new_order" is the one we recognize, - # we return the num_denum of the dimshuffled input. - return self.get_num_denum(inp.owner.inputs[0]) - else: - # This is when the input isn't produced by main, - # inverse or reciprocal. - return [inp], [] + # The graph is walked with an explicit stack of (variable, inverted) + # pairs, where inverted variables accumulate into denum. Children are + # pushed in reverse so elements land in left-to-right traversal order. + # Op dispatch checks identity before (the much slower) equality, since + # graphs are almost always built from the singleton mul/div/... ops. + main, inverse, reciprocal = self.main, self.inverse, self.reciprocal + num, denum = [], [] + stack = [(inp, False)] + while stack: + v, inverted = stack.pop() + apply = v.owner + if apply is None: + (denum if inverted else num).append(v) + continue + op = apply.op + if op is main or op == main: + # main(x, y, ...) contributes each argument's num/denum as is + stack.extend((i, inverted) for i in reversed(apply.inputs)) + elif op is inverse or op == inverse: + # inverse(x, y) contributes y with num/denum flipped + x, y = apply.inputs + stack.append((y, not inverted)) + stack.append((x, inverted)) + elif op is reciprocal or op == reciprocal: + # reciprocal(x) contributes x with num/denum flipped + stack.append((apply.inputs[0], not inverted)) + elif isinstance(op, DimShuffle) and ( + op.is_left_expand_dims + or op.new_order == tuple(range(apply.inputs[0].type.ndim)) + ): + # DimShuffles that only complete the shape with broadcastable + # dimensions to the left are inserted automatically by Elemwise + # when inputs have different numbers of dimensions. They are + # transparent to the walk: merge_num_denum recreates them. + stack.append((apply.inputs[0], inverted)) else: - return [inp], [] - num = [] - denum = [] - parent = inp.owner - - # We get the (num, denum) pairs for each input - # pairs = [self.get_num_denum(input2) if input2.type.dtype == - # input.type.dtype else ([input2], []) for input2 in - # parent.inputs] - pairs = [self.get_num_denum(input2) for input2 in parent.inputs] - - if parent.op == self.main: - # If we have main(x, y, ...), numx, denumx, numy, denumy, ... - # then num is concat(numx, numy, num...) and denum is - # concat(denumx, denumy, denum...) note that main() can have any - # number of arguments >= 0 concat is list concatenation - num = reduce(list.__iadd__, map(operator.itemgetter(0), pairs)) - denum = reduce(list.__iadd__, map(operator.itemgetter(1), pairs)) - elif parent.op == self.inverse: - # If we have inverse(x, y), numx, denumx, numy and denumy - # then num is concat(numx, denumy) and denum is - # concat(denumx, numy) note that inverse() is binary - num = pairs[0][0] + pairs[1][1] - denum = pairs[0][1] + pairs[1][0] - elif parent.op == self.reciprocal: - # If we have reciprocal(x), numx, denumx - # then num is denumx and denum is numx - # note that reciprocal() is unary - num = pairs[0][1] - denum = pairs[0][0] + (denum if inverted else num).append(v) return num, denum def merge_num_denum(self, num, denum): @@ -1369,6 +1348,8 @@ def simplify_factors(self, num, denum): """ ln = len(num) ld = len(denum) + if not (ln and ld): + return num, denum if ld > 2 and ln > 2: # Faster version for "big" inputs. while True: @@ -1426,6 +1407,11 @@ def simplify_constants(self, orig_num, orig_denum, out_type=None): else: denum.append(v) + if not numct and not denumct and (self.use_reciprocal or num): + # No constants to fold; `calculate` would return the neutral + # element and everything below would be a no-op. + return num, denum + if self.use_reciprocal or num: # This will calculate either: # [inverse(main(*numct), main(*denumct))] @@ -1444,6 +1430,19 @@ def simplify_constants(self, orig_num, orig_denum, out_type=None): # Wrapping ct in a Constant with the right dtype ct = [constant(c, dtype=out_type.dtype) for c in ct] + if ( + self.absorbing_element is not None + and len(ct) == 1 + and (num or denum) + and np.all(ct[0].data == self.absorbing_element) + ): + # The constants folded to the absorbing element of `main`, so every + # surviving factor is irrelevant: main(0, x, y) == 0. This has to + # come before the single-constant shortcut below, which would + # otherwise hand back `orig_num` untouched and leave the factors in + # place. + return ct, [] + if orig_num and len(numct) == 1 and len(denumct) == 0 and ct: # In that case we should only have one constant in `ct`. [var_ct] = ct @@ -1475,15 +1474,33 @@ def simplify_constants(self, orig_num, orig_denum, out_type=None): return ct + num, denum + def _broadcast_like_output(self, new, node): + # Equivalent to broadcast_arrays(new, *node.inputs)[0], without + # building the n-1 chains whose results would be discarded. + for inp in node.inputs: + new = second(inp, new) + return new + def transform(self, fgraph, node, enforce_tracks=True): + # Op checks test identity before (the much slower) equality, since + # graphs are almost always built from the singleton mul/div/... ops. op = node.op - if enforce_tracks and (op not in {self.main, self.inverse, self.reciprocal}): + main, inverse, reciprocal = self.main, self.inverse, self.reciprocal + if enforce_tracks and not ( + op is main + or op is inverse + or op is reciprocal + or op == main + or op == inverse + or op == reciprocal + ): return False assert len(node.outputs) == 1 out = node.outputs[0] - out_clients = fgraph.clients.get(out) + clients = fgraph.clients + out_clients = clients.get(out) if not out_clients: return False @@ -1491,27 +1508,80 @@ def transform(self, fgraph, node, enforce_tracks=True): # check if any of the clients of this node would be part of # this canonized graph... if so, we do nothing and wait for # them to be transformed. - for c, c_idx in out_clients: - while ( - isinstance(c.op, DimShuffle) and len(fgraph.clients[c.outputs[0]]) <= 1 + for c, _ in out_clients: + c_op = c.op + while isinstance(c_op, DimShuffle): + c_out_clients = clients[c.outputs[0]] + if len(c_out_clients) > 1: + break + c = c_out_clients[0][0] + c_op = c.op + if ( + c_op is main + or c_op is inverse + or c_op is reciprocal + or c_op == main + or c_op == inverse + or c_op == reciprocal ): - c = fgraph.clients[c.outputs[0]][0][0] - if c.op in [self.main, self.inverse, self.reciprocal]: return False + absorbing = self.absorbing_element + if absorbing is not None and (op is main or op == main): + # main(0, x, y) == 0. Checking the direct inputs lets us skip the + # whole get_num_denum walk for what is by far the common case. + # Restricted to `main`: for `inverse` the element only absorbs in + # the numerator position (0 / x == 0, but x / 0 != 0). + for inp in node.inputs: + if isinstance(inp, TensorConstant): + # The cached unique_value avoids the exception-driven + # constant walk for the direct constant case. + value = inp.unique_value + if value is None or value != absorbing: + continue + elif (inp_node := inp.owner) is not None: + # Inputs produced by main/inverse/reciprocal need no walk: + # any absorbing constant they hide is recovered by the + # get_num_denum flatten and folded in simplify_constants. + inp_op = inp_node.op + if ( + inp_op is main + or inp_op is inverse + or inp_op is reciprocal + or inp_op == main + or inp_op == inverse + or inp_op == reciprocal + ): + continue + try: + value = get_underlying_scalar_constant_value(inp) + except NotScalarConstantError: + continue + if value != absorbing: + continue + else: + continue + out_type = out.type + # Matching the output ndim up front spares Elemwise from + # wrapping the constant in a DimShuffle. + new = constant( + np.full((1,) * out_type.ndim, absorbing, dtype=out_type.dtype) + ) + if new.type.broadcastable != out_type.broadcastable: + new = self._broadcast_like_output(new, node) + copy_stack_trace(out, new) + return [new] + # Here we make the canonical version of the graph around this node # See the documentation of get_num_denum and simplify - orig_num, orig_denum = self.get_num_denum(node.outputs[0]) + orig_num, orig_denum = self.get_num_denum(out) num, denum = self.simplify(list(orig_num), list(orig_denum), out.type) - def same(x, y): - return len(x) == len(y) and all( - np.all(xe == ye) for xe, ye in zip(x, y, strict=True) - ) - if ( - same(orig_num, num) - and same(orig_denum, denum) + # Variable equality is identity, so these are cheap element-wise + # identity checks. + orig_num == num + and orig_denum == denum and # Check to see if we've collapsed some nested ops. not ( @@ -1526,7 +1596,7 @@ def same(x, y): # Do a similar check for the reciprocal op. not ( self.use_reciprocal - and node.op == self.reciprocal + and (op is reciprocal or op == reciprocal) and len(orig_num) == 0 and node.inputs[0].owner and len(node.inputs[0].owner.inputs) < len(orig_denum) @@ -1539,7 +1609,7 @@ def same(x, y): new = cast(new, out.type.dtype) if new.type.broadcastable != out.type.broadcastable: - new = broadcast_arrays(new, *node.inputs)[0] + new = self._broadcast_like_output(new, node) if (new.type.dtype == out.type.dtype) and ( new.type.broadcastable == out.type.broadcastable @@ -1583,7 +1653,7 @@ def mul_calculate(num, denum, aslist=False, out_type=None): local_mul_canonizer = AlgebraicCanonizer( - mul, true_div, reciprocal, mul_calculate, False + mul, true_div, reciprocal, mul_calculate, False, absorbing_element=0 ) register_canonicalize(local_mul_canonizer, "shape_unsafe", name="local_mul_canonizer") @@ -1591,7 +1661,42 @@ def mul_calculate(num, denum, aslist=False, out_type=None): @register_canonicalize @node_rewriter([neg]) def local_neg_to_mul(fgraph, node): - return [mul(np.array(-1, dtype=node.inputs[0].dtype), node.inputs[0])] + """Rewrite neg(x) as a multiplication by -1. + + The naive emission `mul(-1, x)` hands `local_mul_canonizer` a fresh `Mul` + to flatten and re-fold on the next canonicalize iteration. The cases below + emit the already-canonical form directly, so that hand-off never happens. + """ + [x] = node.inputs + + match x.owner_op_and_inputs: + case Elemwise(ps.Neg()), y: + # neg(neg(y)) -> y + return [y] + + case Elemwise(ps.Mul()), *factors: + for i, factor in enumerate(factors): + # neg(mul(c, *rest)) -> mul(-c, *rest): fold the sign into the + # constant rather than adding a factor. Only valid when negating the + # constant in its own dtype is exact under the output dtype: + # floats/complex always are; matching integer dtypes wrap + # consistently (-(c*x) == (-c)*x mod 2^n), but a narrow int constant + # in a wider mul is not (e.g. uint8 2 must become -2.0, not 254). + if isinstance(factor, TensorConstant) and ( + factor.type.numpy_dtype.kind in "fc" or factor.dtype == x.dtype + ): + new_factors = list(factors) + new_factors[i] = constant(-factor.data, dtype=factor.dtype) + return [mul(*new_factors)] + # neg(mul(*xs)) -> mul(-1, *xs), flat rather than mul(-1, mul(*xs)). + return [mul(np.array(-1, dtype=x.dtype), *factors)] + + case (None,) if isinstance(x, TensorConstant) and x.dtype != "bool": + # neg(c) -> -c + return [constant(-x.data, dtype=x.dtype)] + + case _: + return [mul(np.array(-1, dtype=x.dtype), x)] @register_specialize @@ -2261,12 +2366,17 @@ def local_add_neg_to_sub(fgraph, node): return [new_out] -@register_canonicalize @node_rewriter([mul]) def local_mul_zero(fgraph, node): """ - As part of canonicalization, we replace multiplication by zero - with zero. + Replace multiplication by zero with zero. + + Not registered in canonicalize: `local_mul_canonizer` handles the + absorbing element itself (`absorbing_element=0`), and does so in strictly + more cases -- it also folds `mul(2, x, 0.5, 0)`, where the zero only shows + up after the constants are folded together, which the direct-input scan + below misses. Kept for use in rewrite databases that run without the + canonizer. """ otype = node.outputs[0].type @@ -2584,7 +2694,20 @@ def check_for_x_over_absX(numerators, denominators): # TODO: this function should dig/search through dimshuffles # This won't catch a dimshuffled absolute value for den in list(denominators): - if den.owner and den.owner.op == pt_abs and den.owner.inputs[0] in numerators: + # Identity/scalar_op checks reject non-abs ops without paying for the + # props-based Elemwise equality. + if ( + den.owner + and ( + (den_op := den.owner.op) is pt_abs + or ( + isinstance(den_op, Elemwise) + and isinstance(den_op.scalar_op, ps.Abs) + and den_op == pt_abs + ) + ) + and den.owner.inputs[0] in numerators + ): if den.owner.inputs[0].type.dtype.startswith("complex"): # TODO: Make an Op that projects a complex number to # have unit length but projects 0 to 0. That diff --git a/tests/tensor/rewriting/test_math.py b/tests/tensor/rewriting/test_math.py index 81f7cc8186..2b813f38fa 100644 --- a/tests/tensor/rewriting/test_math.py +++ b/tests/tensor/rewriting/test_math.py @@ -107,11 +107,13 @@ from pytensor.tensor.rewriting.math import ( compute_mul, is_1pexp, + local_add_canonizer, local_div_switch_sink, local_grad_log_erfc_neg, local_greedy_distributor, local_mul_canonizer, local_mul_switch_sink, + local_neg_to_mul, local_reduce_chain, local_reduce_join, local_sum_prod_of_mul_or_div, @@ -284,6 +286,45 @@ def test_muldiv(self, e, exp_g): g_rewritten = rewrite_graph(e, custom_rewrite=mul_canonizer) assert equal_computations([g_rewritten], [exp_g]) + @pytest.mark.parametrize( + "e", + [ + # Absorbing element sitting directly among the factors + lambda x, y: pt.mul(0.0, x, y), + # Other constants fold together with the absorbing element + lambda x, y: pt.mul(2.0, x, 0.5, 0.0, y), + # Absorbing element in the numerator of a division + lambda x, y: pt.mul(0.0, x) / y, + # Absorbing element hidden inside a division factor + lambda x, y: x * (0.0 / y), + ], + ids=["direct", "among_other_constants", "numerator", "in_subfactor"], + ) + def test_mul_absorbing_element(self, e): + """`mul` has an absorbing element: mul(0, x, y) == 0, factors and all.""" + x, y = pt.scalars("xy") + result = RewriteTester( + [x, y], [e(x, y)], include=[], custom_rewrite=local_mul_canonizer + ) + result.assert_graph(constant(0.0, dtype=config.floatX)) + + def test_mul_absorbing_element_broadcast(self): + """The absorbed factors still contribute to the shape of the output.""" + x, y = self.x, self.y + zero = constant(np.zeros((1, 1), dtype=config.floatX)) + result = RewriteTester( + [x, y], [pt.mul(zero, x, y)], include=[], custom_rewrite=local_mul_canonizer + ) + result.assert_graph(second(y, second(x, second(zero, zero)))) + + def test_add_has_no_absorbing_element(self): + """`add` only has a neutral element, so 0 must not absorb its terms.""" + x, y = self.x, self.y + result = RewriteTester( + [x, y], [pt.add(0.0, x, y)], include=[], custom_rewrite=local_add_canonizer + ) + result.assert_graph(x + y) + def test_elemwise_multiple_inputs_rewrites(self): """Verify that the `AlgebraicCanonizer` merges sequential ``Elemwise({mul,add})``.""" # Test with and without DimShuffle @@ -1550,6 +1591,78 @@ def test_stacktrace(self): assert check_stack_trace(f, ops_to_check="last") +class TestLocalNegToMul: + """`neg(x)` becomes a multiplication by -1, already in canonical form.""" + + def test_neg_of_neg(self): + x = matrix("x") + assert rewrite_graph(-(-x), include=("canonicalize",)) is x + + def test_neg_of_mul_absorbs_constant(self): + """neg(mul(c, x)) -> mul(-c, x), rather than mul(-1, mul(c, x)).""" + x = matrix("x") + out = rewrite_graph(-(2.0 * x), include=("canonicalize",)) + + assert isinstance(out.owner.op, Elemwise) + assert isinstance(out.owner.op.scalar_op, ps.Mul) + # Just the negated constant and `x` -- no leftover -1 factor. + assert len(out.owner.inputs) == 2 + [const] = [i for i in out.owner.inputs if isinstance(i, TensorConstant)] + assert np.all(const.data == -2.0) + + @pytest.mark.parametrize("dtype, value", [("uint8", 2), ("int8", -128)]) + def test_neg_of_mul_wrapping_int_constant(self, dtype, value): + """The sign is not folded into an int constant narrower than the mul. + + Negating uint8(2) wraps to 254 and negating int8(-128) wraps to -128, + neither of which is the negation under the wider output dtype. + """ + x = pt.vector("x") # float64 + c = pt.constant(np.array([value], dtype=dtype)) + + result = RewriteTester( + [x], [-pt.mul(c, x)], include=None, custom_rewrite=local_neg_to_mul + ) + result.assert_graph(pt.mul(np.array(-1.0), c, x)) + result.assert_eval(np.array([1.0, 2.0])) + + def test_neg_of_mul_same_int_dtype_folds(self): + """Same-dtype integer constants wrap consistently, so the fold fires.""" + x = pt.vector("x", dtype="int8") + c = pt.constant(np.array([3], dtype="int8")) + + result = RewriteTester( + [x], [-pt.mul(c, x)], include=None, custom_rewrite=local_neg_to_mul + ) + result.assert_graph(pt.mul(pt.constant(np.array([-3], dtype="int8")), x)) + result.assert_eval(np.array([1, 2], dtype="int8")) + + def test_neg_of_mul_stays_flat(self): + """neg(mul(x, y)) -> mul(-1, x, y), rather than mul(-1, mul(x, y)).""" + x, y = matrices("xy") + out = rewrite_graph(-(x * y), include=("canonicalize",)) + + assert isinstance(out.owner.op, Elemwise) + assert isinstance(out.owner.op.scalar_op, ps.Mul) + assert len(out.owner.inputs) == 3 + assert not any( + var.owner is not None + and isinstance(var.owner.op, Elemwise) + and isinstance(var.owner.op.scalar_op, ps.Mul) + for var in out.owner.inputs + ) + + def test_values(self): + x, y = matrices("xy") + xv = np.random.random((3, 3)) + yv = np.random.random((3, 3)) + f = function([x, y], [-(-x), -(2.0 * x), -(x * y)]) + neg_neg, neg_mul_const, neg_mul = f(xv, yv) + np.testing.assert_allclose(neg_neg, xv) + np.testing.assert_allclose(neg_mul_const, -2.0 * xv) + np.testing.assert_allclose(neg_mul, -(xv * yv)) + + def test_local_mul_specialize(): mode = config.mode if mode == "FAST_COMPILE":