Issue passing multiple sequences to aesara.scan

Hello

I’m getting an error when passing multiple sequences to aesara.scan:

import aesara
import pymc as pm
import numpy as np

for module in aesara, pm, np:
    print(module.__name__, module.__version__)

# Generate some dummy data
m = 42
n = 8
g = np.random.randint(1, 5, (m, n))
y = np.random.uniform(size=(m, n))

with pm.Model():
    p = pm.Beta("p", alpha=1, beta=100)
    i = pm.Bernoulli("i", p, size=(m, n)) * 1.0
    rho = pm.Beta("r", alpha=10, beta=1)
    y0 = aesara.tensor.zeros(m)
    mu, _ = aesara.scan(
        lambda i, g, y, rho: rho**g * y + i,
        sequences=[i.T, g.T],
        outputs_info=y0,
        non_sequences=[rho],
    )
    sigma = pm.Exponential("sigma", 1)
    pm.Normal("y_hat", mu=mu.T, sigma=sigma, observed=y)
    pm.sample()
aesara 2.8.7
pymc 4.3.0
numpy 1.23.4
Traceback (most recent call last):
  File "/Users/pattinson/djpiri/content/Notes/2022-10-31-aesara-bug/main.py", line 21, in <module>
    mu, _ = aesara.scan(
  File "/Users/pattinson/.virtualenvs/pymc-fresh-env/lib/python3.10/site-packages/aesara/scan/basic.py", line 681, in scan
    scan_seqs = [seq[:actual_n_steps] for seq in scan_seqs]
  File "/Users/pattinson/.virtualenvs/pymc-fresh-env/lib/python3.10/site-packages/aesara/scan/basic.py", line 681, in <listcomp>
    scan_seqs = [seq[:actual_n_steps] for seq in scan_seqs]
TypeError: slice indices must be integers or None or have an __index__ method

A similar model with only a single sequence (removing g) samples fine:


with pm.Model():
    p = pm.Beta("p", alpha=1, beta=100)
    i = pm.Bernoulli("i", p, size=(m, n)) * 1.0
    rho = pm.Beta("r", alpha=10, beta=1)
    y0 = aesara.tensor.zeros(m)
    mu, _ = aesara.scan(
        lambda i, y, rho: rho * y + i,
        sequences=[i.T],
        outputs_info=y0,
        non_sequences=[rho],
    )
    sigma = pm.Exponential("sigma", 1)
    pm.Normal("y_hat", mu=mu.T, sigma=sigma, observed=y)
    pm.sample()

Is this a problem with aesara? Or have I misunderstood something?

Thanks a lot
David

I believe this is because g is a numpy array, not an aesara tensor. Use the function aesara.tensor.as_tensor_variable to convert it before you pass it into the scan, and things should work.

That’s it, thanks a lot!