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])