pytensor
pytensor copied to clipboard
Optimize `Sum`s of `MakeVector`s and `Join`s
Please describe the purpose of filing this issue
Not a drastic improvement by any means, but something we can keep in mind:
reduce(at.concatenate(*tensors)) -> reduce(reduce(tensor) for tensor in tensors)
Ignoring any axis complexities
import pytensor
import pytensor.tensor as pt
import numpy as np
x = pt.vector("x")
y = pt.vector("y")
f1 = pytensor.function([x, y], pt.sum(pt.concatenate((x, y))))
f2 = pytensor.function([x, y], pt.sum((pt.sum(x), pt.sum(y))))
f3 = pytensor.function([x, y], pt.add(pt.sum(x), pt.sum(y)))
pytensor.dprint(f1)
print()
pytensor.dprint(f2)
print()
pytensor.dprint(f3)
x_val = np.random.rand(100_000)
y_val = np.random.rand(200_000)
%timeit f1(x_val, y_val)
%timeit f2(x_val, y_val)
%timeit f3(x_val, y_val)
Sum{acc_dtype=float64} [id A] '' 1
|Join [id B] '' 0
|TensorConstant{0} [id C]
|x [id D]
|y [id E]
Sum{acc_dtype=float64} [id A] '' 3
|MakeVector{dtype='float64'} [id B] '' 2
|Sum{acc_dtype=float64} [id C] '' 1
| |x [id D]
|Sum{acc_dtype=float64} [id E] '' 0
|y [id F]
Elemwise{Add}[(0, 0)] [id A] '' 2
|Sum{acc_dtype=float64} [id B] '' 1
| |x [id C]
|Sum{acc_dtype=float64} [id D] '' 0
|y [id E]
544 µs ± 27.5 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)
270 µs ± 5.11 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)
270 µs ± 8.86 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)