Reduction rewrites discard explicit accumulator dtype and intermediate narrowing

Problem

Several reduction rewrites discard an explicitly selected acc_dtype or an intermediate narrowing boundary. This is independent of the PyTorch lowering problem in #2516.

Reproduced with PyTensor 3.3.2 / NumPy 2.5.3 on Apple Silicon. The script below imports only upstream PyTensor and NumPy. Unoptimized and six-exclusion runs produce the expected values.

Rewrite / case	Linker	Expected	FAST_RUN
`local_sum_prod_all_to_none`: uint16 `[65535, 1]`, axis=0, output=float64, accumulator=uint16	Python	0	65536
`local_useless_reduce`: int64 `[16777217]`, axis=(), output=int64, accumulator=float32	Python	16777216	16777217
`local_reduce_broadcastable`: same value, shape `(1, None)`, axis=0	Python	16777216	16777217
`local_reduce_chain`: inner uint16 reduction followed by float64 reduction	Python	0	65536
`local_reduce_join`: float32 vectors `[1e8]`, `[1]`, `[-1e8]`, stacked and summed with float64 accumulator	C	1	0
`local_careduce_of_alloc`: repeat uint16 65535 twice, output=float64, accumulator=uint16	Python	65534	131070

The Join case happens to return 1 with the Python linker; the C linker demonstrates the lost precision boundary. It also fails with our native PyTorch execution of the rewritten graph. This report does not claim that every linker fails in every row.

The all-axes rewrite reconstructs the reduction without acc_dtype; empty/broadcast axes can remove a required cast through the accumulator; chain/join/Alloc transformations can remove a required rounding or modular-overflow step. The current main local_careduce_of_alloc still explicitly says to ignore acc_dtype when replacing a reduction by multiplication; that is not valid for the unsigned example.

Reproducer

Save as repro.py and run python repro.py --output results.json (a working C compiler is needed for the three Join C-linker runs).

import argparse
import json
from pathlib import Path

import numpy as np
import pytensor
import pytensor.tensor as pt
from pytensor.compile.mode import Mode
from pytensor.tensor.type import TensorType

EXCLUDED = (
    "local_sum_prod_all_to_none",
    "local_useless_reduce",
    "local_reduce_broadcastable",
    "local_reduce_chain",
    "local_reduce_join",
    "local_careduce_of_alloc",
)
MODES = {
    "unoptimized_python": Mode(linker="py", optimizer=None),
    "fast_run_python": Mode(linker="py", optimizer="fast_run"),
    "six_exclusions_python": Mode(linker="py", optimizer="fast_run").excluding(*EXCLUDED),
}
parser = argparse.ArgumentParser()
parser.add_argument("--output", required=True)
args = parser.parse_args()
records = []


def run(name, inputs, output, values, expected):
    modes = dict(MODES)
    if name == "joined_reduction_accumulator":
        modes.update(
            unoptimized_c=Mode(linker="c", optimizer=None),
            fast_run_c=Mode(linker="c", optimizer="fast_run"),
            six_exclusions_c=Mode(linker="c", optimizer="fast_run").excluding(*EXCLUDED),
        )
    for label, mode in modes.items():
        fn = pytensor.function(inputs, output, mode=mode)
        record = {
            "case": name,
            "mode": label,
            "expected": np.asarray(expected).tolist(),
            "expected_dtype": str(np.asarray(expected).dtype),
            "optimized_ops": [str(node.op) for node in fn.maker.fgraph.toposort()],
        }
        try:
            actual = np.asarray(fn(*values))
            record.update(
                actual=actual.tolist(),
                actual_dtype=str(actual.dtype),
                matches=bool(actual.dtype == np.asarray(expected).dtype and np.array_equal(actual, expected)),
            )
        except Exception as exc:
            record.update(error_type=type(exc).__name__, error=str(exc), matches=False)
        records.append(record)


x = pt.vector("u16", dtype="uint16")
run("all_axes_accumulator", [x], pt.sum(x, axis=0, dtype="float64", acc_dtype="uint16"), [np.array([65535, 1], dtype="uint16")], np.array(0, dtype="float64"))
x = pt.lvector("integer")
run("empty_axes_accumulator", [x], pt.sum(x, axis=(), dtype="int64", acc_dtype="float32"), [np.array([16777217], dtype="int64")], np.array([16777216], dtype="int64"))
x = TensorType("int64", shape=(1, None))("broadcast")
run("broadcast_axis_accumulator", [x], pt.sum(x, axis=0, dtype="int64", acc_dtype="float32"), [np.array([[16777217]], dtype="int64")], np.array([16777216], dtype="int64"))
x = pt.matrix("chain", dtype="uint16")
run("nested_reduction_dtype", [x], pt.sum(pt.sum(x, axis=0, dtype="uint16", acc_dtype="uint16"), dtype="float64", acc_dtype="float64"), [np.array([[65535, 0], [1, 0]], dtype="uint16")], np.array(0, dtype="float64"))
x, y, z = [pt.fvector(name) for name in ("x", "y", "z")]
run("joined_reduction_accumulator", [x, y, z], pt.sum(pt.stack([x, y, z]), axis=0, dtype="float32", acc_dtype="float64"), [np.array([value], dtype="float32") for value in (1e8, 1, -1e8)], np.array([1], dtype="float32"))
x = pt.scalar("allocated", dtype="uint16")
run("allocated_reduction_accumulator", [x], pt.sum(pt.alloc(x, 2), dtype="float64", acc_dtype="uint16"), [np.array(65535, dtype="uint16")], np.array(65534, dtype="float64"))

result = {"pytensor": pytensor.__version__, "numpy": np.__version__, "excluded": EXCLUDED, "records": records}
Path(args.output).write_text(json.dumps(result, indent=2) + "\n")
print(json.dumps(result), flush=True)
assert all(item["matches"] for item in records if item["mode"] not in ("fast_run_python", "fast_run_c"))

Downstream handling

We currently exclude these six rewrites in our native PyTorch mode, including nested OpFromGraph/Scan compilation, rather than changing global registrations or installed dependencies. Tracking downstream: https://github.com/pymc-labs/pytensor.cpp/issues/100. We can remove the exclusions when an upstream release preserves these explicit precision boundaries and passes these cases.

May I ask why do you care though?

Further isolated reproductions identify five more precision-erasing rewrites in3.3.2:

Rewrite	Linker comparison	Required result	Optimized result
`local_sum_make_vector`	Python unoptimized vs FAST_RUN	[16777216,0]	[16777217,0]
`local_subtensor_of_reduce`	CVM unoptimized vs FAST_RUN	16777216	16777217
`local_sum_prod_of_mul_or_div`	CVM unoptimized vs FAST_RUN	0	-256
`local_sum_of_neg_to_neg_of_sum`	CVM unoptimized vs FAST_RUN	-128	128
`local_sumsqr2dot`	CVM unoptimized vs FAST_RUN	[0]	[256]

The integer cases preserve real modular/narrowing boundaries; they do not depend on floating-point summation order. All five also failed through our native production mode before the exclusions and now pass, including nested-graph regressions.

Four upstream-only CVM cases (extracted from the executed diagnostic):

import json
import numpy as np
import pytensor
import pytensor.tensor as pt
from pytensor.compile.mode import Mode

modes = {'unoptimized_upstream': Mode(linker='cvm', optimizer=None), 'optimized_upstream': Mode(linker='cvm', optimizer='fast_run')}
records = []
def run(name, inputs, output, values, expected):
    for label, mode in modes.items():
        record = {'name': name, 'mode': label, 'expected': np.asarray(expected).tolist(), 'expected_dtype': str(np.asarray(expected).dtype)}
        try:
            fn = pytensor.function(inputs, output, mode=mode)
            record['optimized_ops'] = [str(node.op) for node in fn.maker.fgraph.toposort()]
            actual = np.asarray(fn(*values))
            record.update(actual=actual.tolist(), actual_dtype=str(actual.dtype), matches=bool(actual.dtype == np.asarray(expected).dtype and np.array_equal(actual, expected)))
        except Exception as exc:
            record.update(error_type=type(exc).__name__, error=str(exc), matches=False)
        records.append(record)

m = pt.lmatrix('m')
run('subtensor', [m], pt.sum(m, axis=1, dtype='int64', acc_dtype='float32')[0], [np.array([[16777217]], dtype='int64')], np.array(16777216, dtype='int64'))
x = pt.vector('x', dtype='int8')
scale = pt.scalar('scale', dtype='int8')
run('factored_multiply', [x, scale], pt.sum(x * scale, dtype='int64', acc_dtype='int8'), [np.array([64, 64], dtype='int8'), np.int8(2)], np.array(0, dtype='int64'))
run('moved_negation', [x], pt.sum(-x, dtype='int64', acc_dtype='int8'), [np.array([-128], dtype='int8')], np.array(-128, dtype='int64'))
w, g = pt.matrix('w', dtype='int8'), pt.matrix('g', dtype='int8')
run('sumsqr_to_dot', [w, g], pt.sum(pt.sqr(w.dimshuffle('x', 0, 1) * g.dimshuffle(0, 'x', 1)), axis=(1, 2), dtype='int64', acc_dtype='int8'), [np.array([[4]], dtype='int8'), np.array([[4]], dtype='int8')], np.array([0], dtype='int64'))
print(json.dumps(records, indent=2))

For MakeVector, the C reduction implementation independently misses the empty-axis accumulator cast even without optimization. It is not a valid oracle for this case. The Python-linker comparison isolates the rewrite: unoptimized/protected Python return [16777216,0], ordinary FAST_RUN returns [16777217,0]. All three CVM variants still return the latter, exposing a separate C implementation issue.

import json
import numpy as np
import pytensor
import pytensor.tensor as pt
from pytensor.compile.mode import Mode

a, b = pt.lscalar('a'), pt.lscalar('b')
output = pt.sum(pt.stack([a, b]), axis=(), dtype='int64', acc_dtype='float32')
records = []
for linker in ('py', 'cvm'):
    modes = {'unoptimized': Mode(linker=linker, optimizer=None), 'fast_run': Mode(linker=linker, optimizer='fast_run'), 'protected': Mode(linker=linker, optimizer='fast_run').excluding('local_useless_reduce', 'local_reduce_broadcastable', 'local_sum_make_vector')}
    for label, mode in modes.items():
        fn = pytensor.function([a, b], output, mode=mode)
        actual = fn(np.int64(16777217), np.int64(0))
        records.append({'linker': linker, 'mode': label, 'actual': actual.tolist(), 'dtype': str(actual.dtype), 'expected': [16777216, 0], 'optimized_ops': [str(node.op) for node in fn.maker.fgraph.toposort()]})
print(json.dumps(records, indent=2))

Our local exclusion list now also includes these five names. The inherited PyTorch reduction handlers have separate ownership/accumulator defects, tracked at #2521; excluding rewrites alone does not repair those handlers.

Navigation Menu

Repository navigation

Reduction rewrites discard explicit accumulator dtype and intermediate narrowing #2518

Problem

Reproducer

Downstream handling

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Reduction rewrites discard explicit accumulator dtype and intermediate narrowing #2518

Description

Problem

Reproducer

Downstream handling

Activity

ricardoV94 commented on Oct 19, 2026

cetagostini commented on Oct 19, 2026

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions