If you like the PowerTransformed layer, you can make a function that takes a pymc distribution and returns a customdist:
import pymc as pm
def PowerTransformed(name, in_dist, c, **kwargs):
# Mini pytensor lesson:
# All symbolic variables are the outputs of a computational graph (unless they are root nodes!)
# Each output (that is, thing you compute) is a the result of an Operation, called an "Op".
# "Op"s, along with inputs and outputs that go into that op, are owned by "Apply" nodes.
# Nodes know what are the inputs and outputs to the Op.
# in_dist is itself an output. We can ask for the owner, which is an Apply. The Apply, in turn, knows
# what is the Op (the function that created in_dist in the first place), as well as the inputs
# that were used to compute it (the parameter you passed when you make it, plus some internal metadata stuff)
def in_dist_power(*args, **kwargs):
# Get the Op and use it to build a new distribution of the same type (with whatever new args/kwargs),
# but also raise it to c
return in_dist.owner.op(*args, **kwargs) ** c
# Get the inputs to the Op that computed `in_dist`. In pytensor random variables, the first two inputs are
# always a random generator and a size. We discard these -- CustomDist will make new ones.
rng, size, *inputs = in_dist.owner.inputs
# This is totally unnecessary, but illustrates what the owner (the Apply) does: it holds
# *all* information about a specific piece of computation.
# An output is itself a part of that computation, so we get loops of ownership like this:
assert in_dist in in_dist.owner.outputs # True
# Make a new PyMC variable
return pm.CustomDist(name,
*inputs,
dist=in_dist_power,
**kwargs)
Usage:
with pm.Model() as m:
alpha = pm.Exponential('alpha')
beta = pm.Exponential('beta')
ig = pm.InverseGamma.dist(alpha, beta)
x = PowerTransformed('sqrt_inv_gamma', ig, c=0.5)
# Since we forwarded the arguments inside PowerTransformed, the resulting RV will still have
# whatever we want as inputs. So `x` still depends on alpha and beta!
assert all(rv in x.owner.inputs for rv in [alpha, beta])
# Sample the prior using NUTS to show that we can evalute the logp.
# Obviously it would be easier to just call sample_prior_predictive
idata = pm.sample()
For more info about how all this works under the hood, you can read our docs about graph structures.