numba max reduction is 10x slower than sum on the same array

[jessegrabowski]

pt.max reduces at about a tenth the speed of pt.sum through the same CAReduce codegen on the same contiguous array, and unlike sum it loses to numpy rather than beating it. Every axis configuration I tried, full and partial.

import time

import numpy as np
import pytensor
import pytensor.tensor as pt

x_np = np.random.default_rng(0).normal(size=(200, 200, 200))
x = pt.tensor("x", shape=x_np.shape)


def bench(fn, *args):
    fn(*args)
    start = time.perf_counter()
    for _ in range(10):
        fn(*args)
    return (time.perf_counter() - start) / 10 * 1000


for name, op in (("max", pt.max), ("sum", pt.sum)):
    print(name, round(bench(pytensor.function([x], op(x), trust_input=True), x_np), 2), "ms")

print("numpy max", round(bench(x_np.max), 2), "ms")
print("numpy sum", round(bench(x_np.sum), 2), "ms")

# max 7.31 ms   sum 0.72 ms   numpy max 0.72 ms   numpy sum 0.96 ms

[ricardoV94]

I suspect some missing SIMD. Scalar max is literally keeping the largest of the two operands.

[ricardoV94]

I reproduce this on my system as well