def ar_logp(op, values, rhos, sigma, init_dist, steps, noise_rng, **kwargs):
(value,) = values
ar_order = op.ar_order
constant_term = op.constant_term
# Convolve rhos with values
if constant_term:
expectation = at.add(
rhos[..., 0, None],
*(
rhos[..., i + 1, None] * value[..., ar_order - (i + 1) : -(i + 1)]
for i in range(ar_order)
),
)
else:
expectation = at.add(
*(
rhos[..., i, None] * value[..., ar_order - (i + 1) : -(i + 1)]
for i in range(ar_order)
)
)
# Compute and collapse logp across time dimension
innov_logp = at.sum(
logp(Normal.dist(0, sigma[..., None]), value[..., ar_order:] - expectation), axis=-1
)
init_logp = logp(init_dist, value[..., :ar_order])
if init_dist.owner.op.ndim_supp == 0:
init_logp = at.sum(init_logp, axis=-1)
return init_logp + innov_logp
def test_clone_new_inputs():
"""Make sure that `Apply.clone_with_new_inputs` properly handles `Type` changes."""
x = at.tensor(np.float64, shape=(None,))
y = at.tensor(np.float64, shape=(1,))
z = at.add(x, y)
assert z.type.shape == (None,)
x_new = at.tensor(np.float64, shape=(1,))
# The output nodes should be reconstructed, because the input types' static
# shape information increased in specificity
z_node_new = z.owner.clone_with_new_inputs([x_new, y])
assert z_node_new.outputs[0].type.shape == (1,)
assert z_node_new.inputs[0].type.shape == (1,)
assert z_node_new.inputs[1].type.shape == (1,)
# Now, attempt to decrease the specificity of the first input's static
# shape information, but, because we're using strict conversion, we
# shouldn't lose any information
z = at.add(x_new, y)
assert z.type.shape == (1,)
z_node_new = z.owner.clone_with_new_inputs([x, y], strict=True)
assert z_node_new.outputs[0].type.shape == (1,)
assert z_node_new.inputs[0].type.shape == (1,)
assert z_node_new.inputs[1].type.shape == (1,)
def logp(self, value):
"""
Calculate log-probability of AR distribution at specified value.
Parameters
----------
value: numeric
Value for which log-probability is calculated.
Returns
-------
TensorVariable
"""
if self.constant:
x = at.add(*[
self.rho[i + 1] * value[self.p - (i + 1):-(i + 1)]
for i in range(self.p)
])
eps = value[self.p:] - self.rho[0] - x
else:
if self.p == 1:
x = self.rho * value[:-1]
else:
x = at.add(*[
self.rho[i] * value[self.p - (i + 1):-(i + 1)]
for i in range(self.p)
])
eps = value[self.p:] - x
innov_like = Normal.dist(mu=0.0, tau=self.tau).logp(eps)
init_like = self.init.logp(value[:self.p])
return at.sum(innov_like) + at.sum(init_like)
def sum_logdets(self):
dets = [self.logdet]
current = self
while not current.isroot:
current = current.parent
dets.append(current.logdet)
return aet.add(*dets)
def __call__(self, X):
XY = X.dot(X.T)
x2 = at.sum(X**2, axis=1).dimshuffle(0, "x")
X2e = at.repeat(x2, X.shape[0], axis=1)
H = X2e + X2e.T - 2.0 * XY
V = at.sort(H.flatten())
length = V.shape[0]
# median distance
m = at.switch(
at.eq((length % 2), 0),
# if even vector
at.mean(V[((length // 2) - 1):((length // 2) + 1)]),
# if odd vector
V[length // 2],
)
h = 0.5 * m / at.log(floatX(H.shape[0]) + floatX(1))
# RBF
Kxy = at.exp(-H / h / 2.0)
# Derivative
dxkxy = -at.dot(Kxy, X)
sumkxy = at.sum(Kxy, axis=-1, keepdims=True)
dxkxy = at.add(dxkxy, at.mul(X, sumkxy)) / h
return Kxy, dxkxy
def elemwise_logp(model, var):
terms = []
for v in model.basic_RVs:
v_logp = logpt(v)
if var in graph_inputs([v_logp]):
terms.append(v_logp)
return model.fn(at.add(*terms))
def test_joint_logp_subtensor():
"""Make sure we can compute a log-likelihood for ``Y[I]`` where ``Y`` and ``I`` are random variables."""
size = 5
mu_base = floatX(np.power(10, np.arange(np.prod(size)))).reshape(size)
mu = np.stack([mu_base, -mu_base])
sigma = 0.001
rng = aesara.shared(np.random.RandomState(232), borrow=True)
A_rv = Normal.dist(mu, sigma, rng=rng)
A_rv.name = "A"
p = 0.5
I_rv = Bernoulli.dist(p, size=size, rng=rng)
I_rv.name = "I"
A_idx = A_rv[I_rv, at.ogrid[A_rv.shape[-1]:]]
assert isinstance(A_idx.owner.op,
(Subtensor, AdvancedSubtensor, AdvancedSubtensor1))
A_idx_value_var = A_idx.type()
A_idx_value_var.name = "A_idx_value"
I_value_var = I_rv.type()
I_value_var.name = "I_value"
A_idx_logps = joint_logp(A_idx, {
A_idx: A_idx_value_var,
I_rv: I_value_var
},
sum=False)
A_idx_logp = at.add(*A_idx_logps)
logp_vals_fn = aesara.function([A_idx_value_var, I_value_var], A_idx_logp)
# The compiled graph should not contain any `RandomVariables`
assert_no_rvs(logp_vals_fn.maker.fgraph.outputs[0])
decimals = select_by_precision(float64=6, float32=4)
for i in range(10):
bern_sp = sp.bernoulli(p)
I_value = bern_sp.rvs(size=size).astype(I_rv.dtype)
norm_sp = sp.norm(mu[I_value, np.ogrid[mu.shape[1]:]], sigma)
A_idx_value = norm_sp.rvs().astype(A_idx.dtype)
exp_obs_logps = norm_sp.logpdf(A_idx_value)
exp_obs_logps += bern_sp.logpmf(I_value)
logp_vals = logp_vals_fn(A_idx_value, I_value)
np.testing.assert_almost_equal(logp_vals,
exp_obs_logps,
decimal=decimals)
def test_compile_pymc_custom_update_op(self):
"""Test that custom MeasurableVariable Op updates are used by compile_pymc"""
class UnmeasurableOp(OpFromGraph):
def update(self, node):
return {node.inputs[0]: node.inputs[0] + 1}
dummy_inputs = [at.scalar(), at.scalar()]
dummy_outputs = [at.add(*dummy_inputs)]
dummy_x = UnmeasurableOp(dummy_inputs,
dummy_outputs)(aesara.shared(1.0), 1.0)
# Check that there are no updates at first
fn = compile_pymc(inputs=[], outputs=dummy_x)
assert fn() == fn() == 2.0
# And they are enabled once the Op is registered as Measurable
MeasurableVariable.register(UnmeasurableOp)
fn = compile_pymc(inputs=[], outputs=dummy_x)
assert fn() == 2.0
assert fn() == 3.0
def logpt(
var: TensorVariable,
rv_values: Optional[Union[TensorVariable, Dict[TensorVariable,
TensorVariable]]] = None,
*,
jacobian: bool = True,
scaling: bool = True,
transformed: bool = True,
sum: bool = True,
**kwargs,
) -> TensorVariable:
"""Create a measure-space (i.e. log-likelihood) graph for a random variable
or a list of random variables at a given point.
The input `var` determines which log-likelihood graph is used and
`rv_value` is that graph's input parameter. For example, if `var` is
the output of a ``NormalRV`` ``Op``, then the output is a graph of the
density function for `var` set to the value `rv_value`.
Parameters
==========
var
The `RandomVariable` output that determines the log-likelihood graph.
Can also be a list of variables. The final log-likelihood graph will
be the sum total of all individual log-likelihood graphs of variables
in the list.
rv_values
A variable, or ``dict`` of variables, that represents the value of
`var` in its log-likelihood. If no `rv_value` is provided,
``var.tag.value_var`` will be checked and, when available, used.
jacobian
Whether or not to include the Jacobian term.
scaling
A scaling term to apply to the generated log-likelihood graph.
transformed
Apply transforms.
sum
Sum the log-likelihood.
"""
# TODO: In future when we drop support for tag.value_var most of the following
# logic can be removed and logpt can just be a wrapper function that calls aeppl's
# joint_logprob directly.
# If var is not a list make it one.
if not isinstance(var, list):
var = [var]
# If logpt isn't provided values and the variable (provided in var)
# is an RV, it is assumed that the tagged value var or observation is
# the value variable for that particular RV.
if rv_values is None:
rv_values = {}
for _var in var:
if isinstance(_var.owner.op, RandomVariable):
rv_value_var = getattr(_var.tag, "observations",
getattr(_var.tag, "value_var", _var))
rv_values = {_var: rv_value_var}
elif not isinstance(rv_values, Mapping):
# Else if we're given a single value and a single variable we assume a mapping among them.
rv_values = ({
var[0]:
at.as_tensor_variable(rv_values).astype(var[0].type)
} if len(var) == 1 else {})
# Since the filtering of logp graph is based on value variables
# provided to this function
if not rv_values:
warnings.warn(
"No value variables provided the logp will be an empty graph")
if scaling:
rv_scalings = {}
for _var in var:
rv_value_var = getattr(_var.tag, "observations",
getattr(_var.tag, "value_var", _var))
rv_scalings[rv_value_var] = _get_scaling(
getattr(_var.tag, "total_size", None), rv_value_var.shape,
rv_value_var.ndim)
# Aeppl needs all rv-values pairs, not just that of the requested var.
# Hence we iterate through the graph to collect them.
tmp_rvs_to_values = rv_values.copy()
transform_map = {}
for node in io_toposort(graph_inputs(var), var):
try:
curr_vars = [node.default_output()]
except ValueError:
curr_vars = node.outputs
for curr_var in curr_vars:
rv_value_var = getattr(curr_var.tag, "observations",
getattr(curr_var.tag, "value_var", None))
if rv_value_var is None:
continue
rv_value = rv_values.get(curr_var, rv_value_var)
tmp_rvs_to_values[curr_var] = rv_value
# Along with value variables we also check for transforms if any.
if hasattr(rv_value_var.tag, "transform") and transformed:
transform_map[rv_value] = rv_value_var.tag.transform
transform_opt = TransformValuesOpt(transform_map)
temp_logp_var_dict = factorized_joint_logprob(tmp_rvs_to_values,
extra_rewrites=transform_opt,
use_jacobian=jacobian,
**kwargs)
# aeppl returns the logpt for every single value term we provided to it. This includes
# the extra values we plugged in above so we need to filter those out.
logp_var_dict = {}
for value_var, _logp in temp_logp_var_dict.items():
if value_var in rv_values.values():
logp_var_dict[value_var] = _logp
# If it's an empty dictionary the logp is None
if not logp_var_dict:
logp_var = None
else:
# Otherwise apply appropriate scalings and at.add and/or at.sum the
# graphs accordingly.
if scaling:
for _value in logp_var_dict.keys():
if _value in rv_scalings:
logp_var_dict[_value] *= rv_scalings[_value]
if len(logp_var_dict) == 1:
logp_var_dict = tuple(logp_var_dict.values())[0]
if sum:
logp_var = at.sum(logp_var_dict)
else:
logp_var = logp_var_dict
else:
if sum:
logp_var = at.sum(
[at.sum(factor) for factor in logp_var_dict.values()])
else:
logp_var = at.add(*logp_var_dict.values())
# Recompute test values for the changes introduced by the replacements
# above.
if config.compute_test_value != "off":
for node in io_toposort(graph_inputs((logp_var, )), (logp_var, )):
compute_test_value(node)
return logp_var
def infer_shape(self, fgraph, nodes, shapes):
first, second = zip(*shapes)
return [(at.add(*first), at.add(*second))]
def __call__(self, X):
return at.add(self.m1(X), self.m2(X))
def elemwise_logp(model, var):
terms = [
v.logp_elemwiset for v in model.basic_RVs
if var in graph_inputs([v.logpt])
]
return model.fn(add(*terms))
