Why does this scan operation not work?

[justindomke]

Hello all! I was interested in trying out aeppl, and came up with the following extremely simple model representing a gaussian random walk:

import aesara, aeppl

srng = aesara.tensor.random.utils.RandomStream(seed=234)

X0 = srng.normal(0.0, 1)

def fn(x_last):
    return srng.normal(x_last,1.0)
X,_ = aesara.scan(fn,outputs_info=[X0],strict=True,n_steps=10)

I can sample from this model using

f = aesara.function([],[X0,X])
f()

and everything looks good. However, if I try to use aeppl to get the joint log probability by doing

aeppl.joint_logprob(X0,X)

I get a very dramatic error message ending with

RuntimeError: The logprob terms of the following random variables could not be derived: {Subtensor{int64::}.0}

Any idea what I might be doing wrong here?

[brandonwillard]

I believe this is related to #148.

I can take a look at it again shortly, if there's an active interest in this functionality.

[justindomke]

Thanks for the reply. Oddly, the refactored version doesn't run for me either. I've tried the latest stable version and the development version of aeppl. I also created a fresh conda environment (python 3.11) and then only installed aeppl. Regardless of what I do, when I get to this line:

logprob, (X_vv, X0_vv) = aeppl.joint_logprob(X, X0)

I get this error:

  File "/python3.11/site-packages/aeppl/scan.py", line 442, in <genexpr>
    fgraph.shape_feature.get_shape(full_out, i) for i in range(full_out.ndim)
    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
TypeError: ShapeFeature.get_shape() missing 1 required positional argument: 'idx'

and this traceback:

File /python3.11/site-packages/aeppl/joint_logprob.py:177, in conditional_logprob(realized, ir_rewriter, extra_rewrites, *random_variables, **kwargs)
    174 outputs = get_measurable_outputs(rv_node.op, rv_node)
    176 if not outputs:
--> 177     raise DensityNotFound(f"Couldn't derive a log-probability for {rv_node}")
    179 if len(outputs) < len(rv_val_pairs):
    180     raise ValueError(
    181         f"Too many values ({rv_val_pairs}) bound to node {rv_node}."
    182     )

DensityNotFound: Couldn't derive a log-probability for Subtensor{int64::}(for{cpu,scan_fn}.0, ScalarConstant{1})

[brandonwillard]

I created #252 to see if this can be reproduced and debugged in CI.

[brandonwillard]

OK, the issue involved some version mismatches. Once #243 passes, I'll merge it and issue a new release.

[brandonwillard]

Done; should be available on PyPI shortly, and conda-forge a while after that.

[justindomke]

Fantastic! I can confirm that this fixes your refactoring and indeed gives the correct answer.

import aesara
import aesara.tensor as at
import aeppl

srng = at.random.utils.RandomStream(seed=234)

X0 = srng.normal(0.0, 1.0, name="X0")

def fn(x_last):
    return srng.normal(x_last, 1.0)

X, _ = aesara.scan(
    fn, outputs_info=[{"initial": X0, "taps": [-1]}], strict=True, n_steps=10
)

logprob, (X_vv, X0_vv) = aeppl.joint_logprob(X, X0)

logprob_fn = aesara.function([X_vv, X0_vv], logprob)

x0 = 1.123
x = np.linspace(2.345,6.789,10)

print(f'{logprob_fn(x, x0)=}')

# check
from scipy.stats import norm
l = norm.logpdf(x0,0.0,1.0)
x_last = x0
for xi in x:
    l += norm.logpdf(xi,x_last,1.0)
    x_last = xi
print(f'{l=}')

logprob_fn(x, x0)=array(-12.58270459)
l=-12.582704587473621

And I can also confirm that my original example now also works, no refactoring needed!

srng = aesara.tensor.random.utils.RandomStream(seed=234)

X0 = srng.normal(0.0, 1)

def fn(x_last):
    return srng.normal(x_last,1.0)
X,_ = aesara.scan(fn,outputs_info=[X0],strict=True,n_steps=10)

f = aesara.function([],[X0,X])
out = f()
print(f'{out=}')

logprob, (X0_vv, X_vv) = aeppl.joint_logprob(X0,X)
logprob_fn = aesara.function([X0_vv,X_vv], logprob)

x0 = 1.123
x = np.linspace(2.345,6.789,10)
l = logprob_fn(x0,x)
print(f'{l=}')

out=[array(-0.53217654), array([-0.7916821 ,  0.60976271,  1.06650553,  1.49055996,  0.85479389,
        0.68713612, -0.96343455, -1.37174506, -2.52785063, -1.57213096])]
l=array(-12.58270459)

[justindomke]

Fixed, thank you!