Implementing rounding (by manual integration) more efficiently

I tried using both pt.vectorize and vectorize_graph. Both versions use

densities = log((densities[3:][::-1] - densities[:3]) * [5/9, 3/9, 1/9])

though I had to do some reshaping for it to work.

The vectorize version (below) required that all inputs are TensorVariable and not TensorSharedVariable. Is there a way to do the switch explicitly? I didn’t find a way how to and did value + pt.shape_padright([-2, -1, 0, 1, 2, 3]) instead. This version was slower than the original one by roughly 10 %.

def _gamma_cdf(mu, sigma, value):
    dist = pm.Gamma.dist(mu=mu, sigma=sigma)
    return exp(pm.logcdf(dist, value))


def _podium_logp(value, mu, sigma):
    cdf = pt.vectorize(_gamma_cdf)
    densities = cdf(mu, sigma, value + pt.shape_padright([-2, -1, 0, 1, 2, 3]))

    densities = (densities[3:][::-1] - densities[:3]) * pt.shape_padright(
        [5 / 9, 3 / 9, 1 / 9]
    )

    return log(sum(densities, 0))

I managed to get the vectorize_graph version working thanks to your advice but it didn’t help either. It was also slower than the original version by 10 % or so.

def _podium_logp(value, mu, sigma):
    dist = pm.Gamma.dist(mu=mu, sigma=sigma)
    shift = pt.scalar("shift")
    density = exp(pm.logcdf(dist, value + shift))

    densities = vectorize_graph(
        density, replace={shift: pt.as_tensor([-2, -1, 0, 1, 2, 3])}
    )
    densities = (densities[3:][::-1] - densities[:3]) * pt.shape_padright(
        [5 / 9, 3 / 9, 1 / 9]
    )

    return log(sum(densities, 0))