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