pytensor icon indicating copy to clipboard operation
pytensor copied to clipboard

Rewrite homogenous repeat vector

Open ricardoV94 opened this issue 2 months ago • 0 comments

Description

If the number of repetitions in pt.repeat is homogeneous, we could replace it by repeat(x, unique_value). Under the hood this ends up as an alloc, skipping the non-C Op altogether

import numpy as np

x = np.arange(600)
r1 = np.array(2, dtype=int)
r2 = np.array([2] * 600, dtype=int)
np.testing.assert_allclose(np.repeat(x, r1), np.repeat(x, r2))
%timeit np.repeat(x, r1)  # 1.32 μs ± 3.21 ns per loop (mean ± std. dev. of 7 runs, 1,000,000 loops each)
%timeit np.repeat(x, r2)  # 1.85 μs ± 3.69 ns per loop (mean ± std. dev. of 7 runs, 1,000,000 loops each)

Not a high priority issue

ricardoV94 avatar Sep 28 '25 11:09 ricardoV94