From 844088d60223be12186d10f95d83287c271c7e43 Mon Sep 17 00:00:00 2001 From: Margus Niitsoo Date: Tue, 18 Aug 2026 10:59:10 +0300 Subject: [PATCH] Numba: combine n-ary Add/Mul terms pairwise Log-depth dependency chains and pairwise-summation accuracy instead of one flat left-to-right chain. Co-Authored-By: Claude Fable 5 --- pytensor/link/numba/dispatch/scalar.py | 19 +++++++++++++++---- tests/link/numba/test_fused_elemwise.py | 9 +++++++++ 2 files changed, 24 insertions(+), 4 deletions(-) diff --git a/pytensor/link/numba/dispatch/scalar.py b/pytensor/link/numba/dispatch/scalar.py index a0662a9c62..dc70b05618 100644 --- a/pytensor/link/numba/dispatch/scalar.py +++ b/pytensor/link/numba/dispatch/scalar.py @@ -1,5 +1,6 @@ import math from hashlib import sha256 +from itertools import zip_longest import numba import numpy as np @@ -190,11 +191,21 @@ def switch(condition, x, y): def binary_to_nary_func(inputs: list[Variable], binary_op_name: str, binary_op: str): - """Create a Numba-compatible N-ary function from a binary function.""" + """Create a Numba-compatible N-ary function from a binary function. + + Terms combine pairwise rather than left-to-right: same associative result, + log-depth dependency chains, and better float accuracy (pairwise summation). + """ var_prefix = "x" if binary_op_name != "x" else "y" input_names = [f"{var_prefix}{i}" for i in range(len(inputs))] input_signature = ", ".join(input_names) - output_expr = binary_op.join(input_names) + terms = input_names + while len(terms) > 1: + terms = [ + f"({a} {binary_op} {b})" if b is not None else a + for a, b in zip_longest(terms[::2], terms[1::2]) + ] + [output_expr] = terms nary_src = f""" def {binary_op_name}({input_signature}): @@ -225,14 +236,14 @@ def pow(x, y): def numba_funcify_Add(op, node, **kwargs): nary_add_fn = binary_to_nary_func(node.inputs, "add", "+") - return numba_basic.numba_njit(nary_add_fn), scalar_op_cache_key(op) + return numba_basic.numba_njit(nary_add_fn), scalar_op_cache_key(op, cache_version=1) @register_funcify_and_cache_key(Mul) def numba_funcify_Mul(op, node, **kwargs): nary_mul_fn = binary_to_nary_func(node.inputs, "mul", "*") - return numba_basic.numba_njit(nary_mul_fn), scalar_op_cache_key(op) + return numba_basic.numba_njit(nary_mul_fn), scalar_op_cache_key(op, cache_version=1) @register_funcify_and_cache_key(Cast) diff --git a/tests/link/numba/test_fused_elemwise.py b/tests/link/numba/test_fused_elemwise.py index f15dd4e1ca..8a16d90073 100644 --- a/tests/link/numba/test_fused_elemwise.py +++ b/tests/link/numba/test_fused_elemwise.py @@ -1080,3 +1080,12 @@ def test_perform_matches(self): np.testing.assert_allclose( perform_fn(xv, yv, idxv)[0], np.sum(xv[idxv] + yv, axis=1), rtol=1e-10 ) + + +def test_wide_add_pairwise_matches(): + rng = np.random.default_rng(0) + xs = [pt.vector(f"x{i}") for i in range(60)] + vals = [rng.normal(size=5) for _ in xs] + + fn = function(xs, sum(xs[1:], xs[0]), mode=get_mode("NUMBA")) + np.testing.assert_allclose(fn(*vals), np.sum(vals, axis=0), rtol=1e-12)