def _get_scaling(total_size, shape, ndim):
"""
Gets scaling constant for logp
Parameters
----------
total_size: int or list[int]
shape: shape
shape to scale
ndim: int
ndim hint
Returns
-------
scalar
"""
if total_size is None:
coef = floatX(1)
elif isinstance(total_size, int):
if ndim >= 1:
denom = shape[0]
else:
denom = 1
coef = floatX(total_size) / floatX(denom)
elif isinstance(total_size, (list, tuple)):
if not all(
isinstance(i, int)
for i in total_size if (i is not Ellipsis and i is not None)):
raise TypeError("Unrecognized `total_size` type, expected "
"int or list of ints, got %r" % total_size)
if Ellipsis in total_size:
sep = total_size.index(Ellipsis)
begin = total_size[:sep]
end = total_size[sep + 1:]
if Ellipsis in end:
raise ValueError(
"Double Ellipsis in `total_size` is restricted, got %r" %
total_size)
else:
begin = total_size
end = []
if (len(begin) + len(end)) > ndim:
raise ValueError("Length of `total_size` is too big, "
"number of scalings is bigger that ndim, got %r" %
total_size)
elif (len(begin) + len(end)) == 0:
return floatX(1)
if len(end) > 0:
shp_end = shape[-len(end):]
else:
shp_end = np.asarray([])
shp_begin = shape[:len(begin)]
begin_coef = [
floatX(t) / shp_begin[i] for i, t in enumerate(begin)
if t is not None
]
end_coef = [
floatX(t) / shp_end[i] for i, t in enumerate(end) if t is not None
]
coefs = begin_coef + end_coef
coef = at.prod(coefs)
else:
raise TypeError(
"Unrecognized `total_size` type, expected int or list of ints, got %r"
% total_size)
return at.as_tensor(floatX(coef))
def test_broadcast_shape():
def shape_tuple(x, use_bcast=True):
if use_bcast:
return tuple(
s if not bcast else 1
for s, bcast in zip(tuple(x.shape), x.broadcastable)
)
else:
return tuple(s for s in tuple(x.shape))
x = np.array([[1], [2], [3]])
y = np.array([4, 5, 6])
b = np.broadcast(x, y)
x_aet = aet.as_tensor_variable(x)
y_aet = aet.as_tensor_variable(y)
b_aet = broadcast_shape(x_aet, y_aet)
assert np.array_equal([z.eval() for z in b_aet], b.shape)
# Now, we try again using shapes as the inputs
#
# This case also confirms that a broadcast dimension will
# broadcast against a non-broadcast dimension when they're
# both symbolic (i.e. we couldn't obtain constant values).
b_aet = broadcast_shape(
shape_tuple(x_aet, use_bcast=False),
shape_tuple(y_aet, use_bcast=False),
arrays_are_shapes=True,
)
assert any(
isinstance(node.op, Assert) for node in applys_between([x_aet, y_aet], b_aet)
)
assert np.array_equal([z.eval() for z in b_aet], b.shape)
b_aet = broadcast_shape(
shape_tuple(x_aet), shape_tuple(y_aet), arrays_are_shapes=True
)
assert np.array_equal([z.eval() for z in b_aet], b.shape)
# These are all constants, so there shouldn't be any asserts in the
# resulting graph.
assert not any(
isinstance(node.op, Assert) for node in applys_between([x_aet, y_aet], b_aet)
)
x = np.array([1, 2, 3])
y = np.array([4, 5, 6])
b = np.broadcast(x, y)
x_aet = aet.as_tensor_variable(x)
y_aet = aet.as_tensor_variable(y)
b_aet = broadcast_shape(x_aet, y_aet)
assert np.array_equal([z.eval() for z in b_aet], b.shape)
b_aet = broadcast_shape(
shape_tuple(x_aet), shape_tuple(y_aet), arrays_are_shapes=True
)
assert np.array_equal([z.eval() for z in b_aet], b.shape)
# TODO: This will work when/if we use a more sophisticated `is_same_graph`
# implementation.
# assert not any(
# isinstance(node.op, Assert)
# for node in graph_ops([x_aet, y_aet], b_aet)
# )
x = np.empty((1, 2, 3))
y = np.array(1)
b = np.broadcast(x, y)
x_aet = aet.as_tensor_variable(x)
y_aet = aet.as_tensor_variable(y)
b_aet = broadcast_shape(x_aet, y_aet)
assert b_aet[0].value == 1
assert np.array_equal([z.eval() for z in b_aet], b.shape)
assert not any(
isinstance(node.op, Assert) for node in applys_between([x_aet, y_aet], b_aet)
)
b_aet = broadcast_shape(
shape_tuple(x_aet), shape_tuple(y_aet), arrays_are_shapes=True
)
assert np.array_equal([z.eval() for z in b_aet], b.shape)
x = np.empty((2, 1, 3))
y = np.empty((2, 1, 1))
b = np.broadcast(x, y)
x_aet = aet.as_tensor_variable(x)
y_aet = aet.as_tensor_variable(y)
b_aet = broadcast_shape(x_aet, y_aet)
assert b_aet[1].value == 1
assert np.array_equal([z.eval() for z in b_aet], b.shape)
# TODO: This will work when/if we use a more sophisticated `is_same_graph`
# implementation.
# assert not any(
# isinstance(node.op, Assert)
# for node in graph_ops([x_aet, y_aet], b_aet)
# )
b_aet = broadcast_shape(
shape_tuple(x_aet), shape_tuple(y_aet), arrays_are_shapes=True
)
assert np.array_equal([z.eval() for z in b_aet], b.shape)
x1_shp_aet = iscalar("x1")
x2_shp_aet = iscalar("x2")
y1_shp_aet = iscalar("y1")
x_shapes = (1, x1_shp_aet, x2_shp_aet)
x_aet = aet.ones(x_shapes)
y_shapes = (y1_shp_aet, 1, x2_shp_aet)
y_aet = aet.ones(y_shapes)
b_aet = broadcast_shape(x_aet, y_aet)
# TODO: This will work when/if we use a more sophisticated `is_same_graph`
# implementation.
# assert not any(
# isinstance(node.op, Assert)
# for node in graph_ops([x_aet, y_aet], b_aet)
# )
res = aet.as_tensor(b_aet).eval(
{
x1_shp_aet: 10,
x2_shp_aet: 4,
y1_shp_aet: 2,
}
)
assert np.array_equal(res, (2, 10, 4))
y_shapes = (y1_shp_aet, 1, y1_shp_aet)
y_aet = aet.ones(y_shapes)
b_aet = broadcast_shape(x_aet, y_aet)
assert isinstance(b_aet[-1].owner.op, Assert)
def logcdfpt(
var: TensorVariable,
rv_values: Optional[Union[TensorVariable, Dict[TensorVariable, TensorVariable]]] = None,
*,
scaling: bool = True,
sum: bool = True,
**kwargs,
) -> TensorVariable:
"""Create a measure-space (i.e. log-cdf) graph for a random variable at a given point.
Parameters
==========
var
The `RandomVariable` output that determines the log-likelihood graph.
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.
"""
if not isinstance(rv_values, Mapping):
rv_values = {var: rv_values} if rv_values is not None else {}
rv_var, rv_value_var = extract_rv_and_value_vars(var)
rv_value = rv_values.get(rv_var, rv_value_var)
if rv_var is not None and rv_value is None:
raise ValueError(f"No value variable specified or associated with {rv_var}")
if rv_value is not None:
rv_value = at.as_tensor(rv_value)
if rv_var is not None:
# Make sure that the value is compatible with the random variable
rv_value = rv_var.type.filter_variable(rv_value.astype(rv_var.dtype))
if rv_value_var is None:
rv_value_var = rv_value
rv_node = rv_var.owner
rng, size, dtype, *dist_params = rv_node.inputs
# Here, we plug the actual random variable into the log-likelihood graph,
# because we want a log-likelihood graph that only contains
# random variables. This is important, because a random variable's
# parameters can contain random variables themselves.
# Ultimately, with a graph containing only random variables and
# "deterministics", we can simply replace all the random variables with
# their value variables and be done.
tmp_rv_values = rv_values.copy()
tmp_rv_values[rv_var] = rv_var
logp_var = _logcdf(rv_node.op, rv_var, tmp_rv_values, *dist_params, **kwargs)
transform = getattr(rv_value_var.tag, "transform", None) if rv_value_var else None
# Replace random variables with their value variables
replacements = rv_values.copy()
replacements.update({rv_var: rv_value, rv_value_var: rv_value})
(logp_var,), _ = rvs_to_value_vars(
(logp_var,),
apply_transforms=False,
initial_replacements=replacements,
)
if sum:
logp_var = at.sum(logp_var)
if scaling:
logp_var *= _get_scaling(
getattr(rv_var.tag, "total_size", None), rv_value.shape, rv_value.ndim
)
# 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)
if rv_var.name is not None:
logp_var.name = f"__logp_{rv_var.name}"
return logp_var
def logpt(
var: TensorVariable,
rv_values: Optional[Union[TensorVariable, Dict[TensorVariable,
TensorVariable]]] = None,
*,
jacobian: bool = True,
scaling: bool = True,
transformed: bool = True,
cdf: bool = False,
sum: bool = False,
**kwargs,
) -> TensorVariable:
"""Create a measure-space (i.e. log-likelihood) graph for a random variable 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.
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.
cdf
Return the log cumulative distribution.
sum
Sum the log-likelihood.
"""
if not isinstance(rv_values, Mapping):
rv_values = {var: rv_values} if rv_values is not None else {}
rv_var, rv_value_var = extract_rv_and_value_vars(var)
rv_value = rv_values.get(rv_var, rv_value_var)
if rv_var is not None and rv_value is None:
raise ValueError(
f"No value variable specified or associated with {rv_var}")
if rv_value is not None:
rv_value = at.as_tensor(rv_value)
if rv_var is not None:
# Make sure that the value is compatible with the random variable
rv_value = rv_var.type.filter_variable(
rv_value.astype(rv_var.dtype))
if rv_value_var is None:
rv_value_var = rv_value
if rv_var is None:
if var.owner is not None:
return _logp(
var.owner.op,
var,
rv_values,
*var.owner.inputs,
jacobian=jacobian,
scaling=scaling,
transformed=transformed,
cdf=cdf,
sum=sum,
)
return at.zeros_like(var)
rv_node = rv_var.owner
rng, size, dtype, *dist_params = rv_node.inputs
# Here, we plug the actual random variable into the log-likelihood graph,
# because we want a log-likelihood graph that only contains
# random variables. This is important, because a random variable's
# parameters can contain random variables themselves.
# Ultimately, with a graph containing only random variables and
# "deterministics", we can simply replace all the random variables with
# their value variables and be done.
tmp_rv_values = rv_values.copy()
tmp_rv_values[rv_var] = rv_var
if not cdf:
logp_var = _logp(rv_node.op, rv_var, tmp_rv_values, *dist_params,
**kwargs)
else:
logp_var = _logcdf(rv_node.op, rv_var, tmp_rv_values, *dist_params,
**kwargs)
transform = getattr(rv_value_var.tag, "transform",
None) if rv_value_var else None
if transform and transformed and not cdf and jacobian:
transformed_jacobian = transform.jacobian_det(rv_var, rv_value)
if transformed_jacobian:
if logp_var.ndim > transformed_jacobian.ndim:
logp_var = logp_var.sum(axis=-1)
logp_var += transformed_jacobian
# Replace random variables with their value variables
replacements = rv_values.copy()
replacements.update({rv_var: rv_value, rv_value_var: rv_value})
(logp_var, ), _ = rvs_to_value_vars(
(logp_var, ),
apply_transforms=transformed and not cdf,
initial_replacements=replacements,
)
if sum:
logp_var = at.sum(logp_var)
if scaling:
logp_var *= _get_scaling(getattr(rv_var.tag, "total_size", None),
rv_value.shape, rv_value.ndim)
# 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)
if rv_var.name is not None:
logp_var.name = "__logp_%s" % rv_var.name
return logp_var