Aesera Scan Altering Distributions

I am noticing some strange behaviour when trying to use aesara.scan with uniform random variables. For example consider this contrived example to randomly add one or two to the previous value in a sequence. (Note I seemed to need the at.ones_like to avoid in the switch statement to avoid a AttributeError: 'float' object has no attribute 'type' error)

def func(switcher, prev):
    offset = pm.math.switch(
        pm.math.lt(switcher, 0.5),
        at.ones_like(prev),
        2 * at.ones_like(prev),
    )
    return prev + offset

with pm.Model() as model:
    switcher = pm.Uniform("switcher", lower=0, upper=1, shape=2)
    result, updates = aesara.scan(
        fn=func,
        outputs_info={"initial": at.ones_like(1.0, dtype="float64")},
        sequences=[switcher],
        n_steps=2,
    )
    pm.Normal("output", mu=result, sigma=0.01)
    trace = pm.sample()

az.plot_trace(trace)

Screenshot 2022-12-07 160015

Why has the distribution over switcher been altered to be either between 0 and 0.5 or 0.5 to 1?? If I comment out the pm.Normal(...) line I get the expected distribution from 0 to 1…what is scan doing to meddle with this?

Ah, I believe this is answered the FAQs. This is not due to use of scan but rather use of switch. Reparametrizing my problem using sigmoid functions seems to work.

1 Like