def moment_censored(op, rv, dist, lower, upper):
moment = at.switch(
at.eq(lower, -np.inf),
at.switch(
at.isinf(upper),
# lower = -inf, upper = inf
0,
# lower = -inf, upper = x
upper - 1,
),
at.switch(
at.eq(upper, np.inf),
# lower = x, upper = inf
lower + 1,
# lower = x, upper = x
(lower + upper) / 2,
),
)
moment = at.full_like(dist, moment)
return moment
def ar_moment(op, rv, rhos, sigma, init_dist, steps, noise_rng):
# Use last entry of init_dist moment as the moment for the whole AR
return at.full_like(rv, moment(init_dist)[..., -1, None])