Need a little guidance on using pytensor.scan()

[mike-lawrence]

How would the following function be implemented using scan?

def f(a):
    b = np.zeros_like(a)
    b[0] = a[0]
    last_a_copied = a[0]
    copy_count = 1
    for i in range(1,a.shape[0]):
        if copy_count>=last_a_copied:
            last_a_copied = a[i]
            copy_count = 0
        b[i] = last_a_copied
        copy_count += 1
    return(b)

a = np.array([2,1,3,2,4,4,3,3,5,5,2],dtype='int64')
b = f(a)
# print(b)  # Should equal: [2 2 3 3 3 4 4 4 4 5 5]

[ricardoV94]

Here is one option

import pytensor
import pytensor.tensor as pt
from pytensor.ifelse import ifelse

def step(a, last_a_copied, copy_count):
    last_a_copied, copy_count = ifelse(
        copy_count >= last_a_copied,
        (a, pt.ones_like(copy_count)),
        (last_a_copied, copy_count + 1)
    )
    return last_a_copied, copy_count

a = pt.vector("a", dtype="int64")
copy_count = a.max() + 1  # Or any value that's larger than any in a
last_a_copied = a[0]
[b, _], _ = pytensor.scan(
    step,
    sequences=[a],
    outputs_info=[last_a_copied, copy_count],
)
b.eval({a: [2,1,3,2,4,4,3,3,5,5,2]})  # array([2, 2, 3, 3, 3, 4, 4, 4, 4, 5, 5])