PyTorch Max/Min/Prod violate empty-axis ownership; Prod ignores accumulator dtype #2521

Problem

The PyTorch Max, Min and Prod handlers return their input tensor when axis=(). These Ops declare no input alias, so callers and inplace analysis can treat the output as independently owned even though changing it changes caller storage.

Reproduced using upstream dispatch directly with PyTensor3.3.2, PyTorch2.14.1 and NumPy2.5.3, CPU; no downstream package import and no torch.compile.

The input should remain [1,2]. Normal FAST_RUN identity elimination may mask this issue; it appears when a valid mode retains the reduction. Our production-mode reproduction confirmed the same caller mutation after excluding precision-erasing identity rewrites.

Prod also ignores acc_dtype: half [[256,256],[1/256,1]], outputhalf/accumulatorfloat64, returns half inf; NumPy accumulation in float64 followed by the declared output conversion returns half256. This occurs with both axis=None and (0,1) in the full graph diagnostic. Intermediate reduction outputs must remain in accumulator precision until all axes have been reduced.

Reproducer

The following is the reduction portion of the executed upstream-only diagnostic. Explicit Op classes are intentional: the high-level pt.min helper can initially build a negated Max graph, so dispatching only its outer node would not exercise Min.

import json
import numpy as np
import pytensor.tensor as pt
import torch
from pytensor.tensor.math import Max, Min, Prod
from pytensor.link.pytorch.dispatch import pytorch_funcify

result = {'empty_reductions': [], 'product': {}}
for reduction in (Max, Min, Prod):
    x = pt.fvector('x')
    output = reduction(axis=())(x)
    assert output.owner.inputs[0] is x
    values = np.array([1, 2], dtype='float32')
    actual = pytorch_funcify(output.owner.op, node=output.owner)(torch.from_numpy(values)).numpy()
    np.testing.assert_array_equal(actual, values)
    shared = np.shares_memory(actual, values)
    actual[0] = 99
    result['empty_reductions'].append({'op': reduction.__name__, 'shares_input_storage': shared, 'input_after_output_mutation': values.tolist(), 'expected_input': [1, 2]})
x = pt.matrix('x', dtype='float16')
output = pt.prod(x, dtype='float16', acc_dtype='float64')
values = np.array([[256, 256], [1 / 256, 1]], dtype='float16')
actual = pytorch_funcify(output.owner.op, node=output.owner)(torch.from_numpy(values)).numpy()
expected = np.prod(values, dtype='float64').astype('float16')
result['product'] = {'actual_repr': repr(actual.item()), 'actual_dtype': str(actual.dtype), 'expected': expected.item(), 'expected_dtype': str(expected.dtype)}

print(json.dumps(result, indent=2))

Downstream handling

We are preserving owned empty-reduction results and accumulator/output conversion locally, including typed Prod, without modifying global dispatch or installed packages. Remove the local bridge when upstream passes empty-axis mutation isolation, declared dtype/accumulator, multi-axis, empty-product and modular-integer cases.

Related Sum lowering issue: #2516. Rewrite precision issue: #2518. Downstream tracking: https://github.com/pymc-labs/pytensor.cpp/issues/100.