def _multinomial_infer_value_domain(cls, **kwargs):
instance = cls.dist_class(**{k: _dummy_tensor(domain) for k, domain in kwargs.items()}, validate_args=False)
return reals(*instance.event_shape)
def __call__(self):
# calls pyro.param so that params are exposed and constraints applied
# should not create any new torch.Tensors after __init__
self.initialize_params()
N_state = self.config["sizes"]["state"]
# initialize gamma to uniform
gamma = Tensor(
torch.zeros((N_state, N_state)),
OrderedDict([("y_prev", bint(N_state))]),
)
N_v = self.config["sizes"]["random"]
N_c = self.config["sizes"]["group"]
log_prob = []
plate_g = Tensor(torch.zeros(N_c), OrderedDict([("g", bint(N_c))]))
# group-level random effects
if self.config["group"]["random"] == "discrete":
# group-level discrete effect
e_g = Variable("e_g", bint(N_v))
e_g_dist = plate_g + dist.Categorical(**self.params["e_g"])(
value=e_g)
log_prob.append(e_g_dist)
eps_g = (plate_g + self.params["eps_g"]["theta"])(e_g=e_g)
elif self.config["group"]["random"] == "continuous":
eps_g = Variable("eps_g", reals(N_state))
eps_g_dist = plate_g + dist.Normal(**self.params["eps_g"])(
value=eps_g)
log_prob.append(eps_g_dist)
else:
eps_g = to_funsor(0.)
N_s = self.config["sizes"]["individual"]
plate_i = Tensor(torch.zeros(N_s), OrderedDict([("i", bint(N_s))]))
# individual-level random effects
if self.config["individual"]["random"] == "discrete":
# individual-level discrete effect
e_i = Variable("e_i", bint(N_v))
e_i_dist = plate_g + plate_i + dist.Categorical(
**self.params["e_i"])(
value=e_i) * self.raggedness_masks["individual"](t=0)
log_prob.append(e_i_dist)
eps_i = (plate_i + plate_g +
self.params["eps_i"]["theta"](e_i=e_i))
elif self.config["individual"]["random"] == "continuous":
eps_i = Variable("eps_i", reals(N_state))
eps_i_dist = plate_g + plate_i + dist.Normal(
**self.params["eps_i"])(value=eps_i)
log_prob.append(eps_i_dist)
else:
eps_i = to_funsor(0.)
# add group-level and individual-level random effects to gamma
gamma = gamma + eps_g + eps_i
N_state = self.config["sizes"]["state"]
# we've accounted for all effects, now actually compute gamma_y
gamma_y = gamma(y_prev="y(t=1)")
y = Variable("y", bint(N_state))
y_dist = plate_g + plate_i + dist.Categorical(
probs=gamma_y.exp() / gamma_y.exp().sum())(value=y)
# observation 1: step size
step_dist = plate_g + plate_i + dist.Gamma(
**{k: v(y_curr=y)
for k, v in self.params["step"].items()})(
value=self.observations["step"])
# step size zero-inflation
if self.config["zeroinflation"]:
step_zi = dist.Categorical(
probs=self.params["zi_step"]["zi_param"](y_curr=y))(
value="zi_step")
step_zi_dist = plate_g + plate_i + dist.Delta(
self.config["MISSING"])(value=self.observations["step"])
step_dist = (step_zi + Stack("zi_step",
(step_dist, step_zi_dist))).reduce(
ops.logaddexp, "zi_step")
# observation 2: step angle
angle_dist = plate_g + plate_i + dist.VonMises(
**{k: v(y_curr=y)
for k, v in self.params["angle"].items()})(
value=self.observations["angle"])
# observation 3: dive activity
omega_dist = plate_g + plate_i + dist.Beta(
**{k: v(y_curr=y)
for k, v in self.params["omega"].items()})(
value=self.observations["omega"])
# dive activity zero-inflation
if self.config["zeroinflation"]:
omega_zi = dist.Categorical(
probs=self.params["zi_omega"]["zi_param"](y_curr=y))(
value="zi_omega")
omega_zi_dist = plate_g + plate_i + dist.Delta(
self.config["MISSING"])(value=self.observations["omega"])
omega_dist = (omega_zi +
Stack("zi_omega",
(omega_dist, omega_zi_dist))).reduce(
ops.logaddexp, "zi_omega")
# finally, construct the term for parallel scan reduction
hmm_factor = step_dist + angle_dist + omega_dist
hmm_factor = hmm_factor * self.raggedness_masks["individual"]
hmm_factor = hmm_factor * self.raggedness_masks["timestep"]
# copy masking behavior of pyro.infer.TraceEnum_ELBO._compute_model_factors
hmm_factor = hmm_factor + y_dist
log_prob.insert(0, hmm_factor)
return log_prob
def forward(self, observations, add_bias=True):
obs_dim = 2 * self.num_sensors
bias_scale = self.log_bias_scale.exp()
obs_noise = self.log_obs_noise.exp()
trans_noise = self.log_trans_noise.exp()
# bias distribution
bias = Variable('bias', reals(obs_dim))
assert not torch.isnan(bias_scale), "bias scales was nan"
bias_dist = dist_to_funsor(
dist.MultivariateNormal(
torch.zeros(obs_dim),
scale_tril=bias_scale *
torch.eye(2 * self.num_sensors)))(value=bias)
init_dist = torch.distributions.MultivariateNormal(torch.zeros(4),
scale_tril=100. *
torch.eye(4))
self.init = dist_to_funsor(init_dist)(value="state")
# hidden states
prev = Variable("prev", reals(4))
curr = Variable("curr", reals(4))
self.trans_dist = f_dist.MultivariateNormal(
loc=prev @ NCV_TRANSITION_MATRIX,
scale_tril=trans_noise * NCV_PROCESS_NOISE.cholesky(),
value=curr)
state = Variable('state', reals(4))
obs = Variable("obs", reals(obs_dim))
observation_matrix = Tensor(
torch.eye(4,
2).unsqueeze(-1).expand(-1, -1,
self.num_sensors).reshape(4, -1))
assert observation_matrix.output.shape == (
4, obs_dim), observation_matrix.output.shape
obs_loc = state @ observation_matrix
if add_bias:
obs_loc += bias
self.observation_dist = f_dist.MultivariateNormal(
loc=obs_loc, scale_tril=obs_noise * torch.eye(obs_dim), value=obs)
logp = bias_dist
curr = "state_init"
logp += self.init(state=curr)
for t, x in enumerate(observations):
prev, curr = curr, f"state_{t}"
logp += self.trans_dist(prev=prev, curr=curr)
logp += self.observation_dist(state=curr, obs=x)
# marginalize out previous state
logp = logp.reduce(ops.logaddexp, prev)
# marginalize out bias variable
logp = logp.reduce(ops.logaddexp, "bias")
# save posterior over the final state
assert set(logp.inputs) == {f'state_{len(observations) - 1}'}
posterior = funsor_to_mvn(logp, ndims=0)
# marginalize out remaining variables
logp = logp.reduce(ops.logaddexp)
assert isinstance(logp, Tensor) and logp.shape == (), logp.pretty()
return logp.data, posterior
eval(expr)
@pytest.mark.parametrize("expr", [
"Variable('x', reals()).log()",
"Number(1) / Variable('x', reals())",
"Variable('x', reals()) ** Number(2)",
"Stack('t', (Number(1), Variable('x', reals()))).reduce(ops.logaddexp, 't')",
])
def test_eager_or_die_error(expr):
with interpretation(eager_or_die):
with pytest.raises(NotImplementedError):
eval(expr)
@pytest.mark.parametrize('domain', [bint(3), reals()])
def test_variable(domain):
x = Variable('x', domain)
check_funsor(x, {'x': domain}, domain)
assert Variable('x', domain) is x
assert x('x') is x
y = Variable('y', domain)
assert x('y') is y
assert x(x='y') is y
assert x(x=y) is y
x4 = Variable('x', bint(4))
assert x4 is not x
assert x4('x') is x4
assert x(y=x4) is x
xp1 = x + 1.
if sampled_vars:
assert dict(y.inputs) == dict(expected_inputs), sampled_vars
else:
assert y is x
@pytest.mark.parametrize('sample_inputs', [
(),
(('s', bint(3)), ),
(('s', bint(3)), ('t', bint(4))),
],
ids=id_from_inputs)
@pytest.mark.parametrize('batch_inputs', [
(),
(('b', bint(2)), ),
(('c', reals()), ),
(('b', bint(2)), ('c', reals())),
],
ids=id_from_inputs)
@pytest.mark.parametrize('event_inputs', [
(('e', reals()), ),
(('e', reals()), ('f', reals(2))),
],
ids=id_from_inputs)
def test_gaussian_shape(sample_inputs, batch_inputs, event_inputs):
be_inputs = OrderedDict(batch_inputs + event_inputs)
expected_inputs = OrderedDict(sample_inputs + batch_inputs + event_inputs)
sample_inputs = OrderedDict(sample_inputs)
batch_inputs = OrderedDict(batch_inputs)
event_inputs = OrderedDict(event_inputs)
x = random_gaussian(be_inputs)
def test_normal_defaults():
loc = Variable('loc', reals())
scale = Variable('scale', reals())
value = Variable('value', reals())
assert dist.Normal(loc, scale) is dist.Normal(loc, scale, value)
def _rand(batch_shape, *event_shape):
inputs = OrderedDict(zip("abcdef", map(bint, reversed(batch_shape))))
return random_tensor(inputs, reals(*event_shape))
def _fill_defaults(v, log_density=0, value='value'):
v = to_funsor(v)
log_density = to_funsor(log_density, reals())
value = to_funsor(value, v.output)
return v, log_density, value
def _fill_defaults(loc, scale, value='value'):
loc = to_funsor(loc, reals())
scale = to_funsor(scale, reals())
value = to_funsor(value, reals())
return loc, scale, value
def _fill_defaults(concentration1, concentration0, value='value'):
concentration1 = to_funsor(concentration1, reals())
concentration0 = to_funsor(concentration0, reals())
value = to_funsor(value, reals())
return concentration1, concentration0, value
def _fill_defaults(total_count, probs, value='value'):
total_count = to_funsor(total_count, reals())
probs = to_funsor(probs)
assert probs.dtype == "real"
value = to_funsor(value, reals())
return total_count, probs, value
def _fill_defaults(logits, value='value'):
logits = to_funsor(logits)
assert logits.dtype == "real"
value = to_funsor(value, reals())
return logits, value
def _fill_defaults(probs, value='value'):
probs = to_funsor(probs)
assert probs.dtype == "real"
value = to_funsor(value, reals())
return probs, value
factors1 = partial_sum_product(sum_op, prod_op, factors, vars1, plates)
factors2 = partial_sum_product(sum_op, prod_op, factors1, vars2, plates)
actual = reduce(prod_op, factors2)
expected = sum_product(sum_op, prod_op, factors, vars1 | vars2, plates)
assert_close(actual, expected)
@pytest.mark.parametrize('num_steps', [None] + list(range(1, 13)))
@pytest.mark.parametrize('sum_op,prod_op,state_domain', [
(ops.add, ops.mul, bint(2)),
(ops.add, ops.mul, bint(3)),
(ops.logaddexp, ops.add, bint(2)),
(ops.logaddexp, ops.add, bint(3)),
(ops.logaddexp, ops.add, reals()),
(ops.logaddexp, ops.add, reals(2)),
],
ids=str)
@pytest.mark.parametrize('batch_inputs', [
{},
{
"foo": bint(5)
},
{
"foo": bint(2),
"bar": bint(4)
},
],
ids=lambda d: ",".join(d.keys()))
@pytest.mark.parametrize('impl', [
def test_beta_density(batch_shape, eager):
batch_dims = ('i', 'j', 'k')[:len(batch_shape)]
inputs = OrderedDict((k, bint(v)) for k, v in zip(batch_dims, batch_shape))
@funsor.torch.function(reals(), reals(), reals(), reals())
def beta(concentration1, concentration0, value):
return torch.distributions.Beta(concentration1,
concentration0).log_prob(value)
check_funsor(beta, {
'concentration1': reals(),
'concentration0': reals(),
'value': reals()
}, reals())
concentration1 = Tensor(torch.randn(batch_shape).exp(), inputs)
concentration0 = Tensor(torch.randn(batch_shape).exp(), inputs)
value = Tensor(torch.rand(batch_shape), inputs)
expected = beta(concentration1, concentration0, value)
check_funsor(expected, inputs, reals())
d = Variable('value', reals())
actual = dist.Beta(concentration1, concentration0, value) if eager else \
dist.Beta(concentration1, concentration0, d)(value=value)
check_funsor(actual, inputs, reals())
assert_close(actual, expected)
def _eager_subs_real(self, subs, remaining_subs):
# Broadcast all component tensors.
subs = OrderedDict(subs)
int_inputs = OrderedDict(
(k, d) for k, d in self.inputs.items() if d.dtype != 'real')
tensors = [
Tensor(self.info_vec, int_inputs),
Tensor(self.precision, int_inputs)
]
tensors.extend(subs.values())
int_inputs, tensors = align_tensors(*tensors)
batch_dim = tensors[0].dim() - 1
batch_shape = broadcast_shape(*(x.shape[:batch_dim] for x in tensors))
(info_vec, precision), values = tensors[:2], tensors[2:]
offsets, event_size = _compute_offsets(self.inputs)
slices = [(k, slice(offset, offset + self.inputs[k].num_elements))
for k, offset in offsets.items()]
# Expand all substituted values.
values = OrderedDict(zip(subs, values))
for k, value in values.items():
value = value.reshape(value.shape[:batch_dim] + (-1, ))
if not torch._C._get_tracing_state():
assert value.size(-1) == self.inputs[k].num_elements
values[k] = value.expand(batch_shape + value.shape[-1:])
# Try to perform a complete substitution of all real variables, resulting in a Tensor.
if all(k in subs for k, d in self.inputs.items() if d.dtype == 'real'):
# Form the concatenated value.
value = BlockVector(batch_shape + (event_size, ))
for k, i in slices:
if k in values:
value[..., i] = values[k]
value = value.as_tensor()
# Evaluate the non-normalized log density.
result = _vv(value, info_vec - 0.5 * _mv(precision, value))
result = Tensor(result, int_inputs)
assert result.output == reals()
return Subs(result, remaining_subs) if remaining_subs else result
# Perform a partial substution of a subset of real variables, resulting in a Joint.
# We split real inputs into two sets: a for the preserved and b for the substituted.
b = frozenset(k for k, v in subs.items())
a = frozenset(k for k, d in self.inputs.items()
if d.dtype == 'real' and k not in b)
prec_aa = torch.cat([
torch.cat([precision[..., i1, i2] for k2, i2 in slices if k2 in a],
dim=-1) for k1, i1 in slices if k1 in a
],
dim=-2)
prec_ab = torch.cat([
torch.cat([precision[..., i1, i2] for k2, i2 in slices if k2 in b],
dim=-1) for k1, i1 in slices if k1 in a
],
dim=-2)
prec_bb = torch.cat([
torch.cat([precision[..., i1, i2] for k2, i2 in slices if k2 in b],
dim=-1) for k1, i1 in slices if k1 in b
],
dim=-2)
info_a = torch.cat([info_vec[..., i] for k, i in slices if k in a],
dim=-1)
info_b = torch.cat([info_vec[..., i] for k, i in slices if k in b],
dim=-1)
value_b = torch.cat([values[k] for k, i in slices if k in b], dim=-1)
info_vec = info_a - _mv(prec_ab, value_b)
log_scale = _vv(value_b, info_b - 0.5 * _mv(prec_bb, value_b))
precision = prec_aa.expand(info_vec.shape + (-1, ))
inputs = int_inputs.copy()
for k, d in self.inputs.items():
if k not in subs:
inputs[k] = d
result = Gaussian(info_vec, precision, inputs) + Tensor(
log_scale, int_inputs)
return Subs(result, remaining_subs) if remaining_subs else result
def test_delta_delta():
v = Variable('v', reals(2))
point = Tensor(torch.randn(2))
log_density = Tensor(torch.tensor(0.5))
d = dist.Delta(point, log_density, v)
assert d is Delta('v', point, log_density)
def _eager_subs_affine(self, subs, remaining_subs):
# Extract an affine representation.
affine = OrderedDict()
for k, v in subs:
const, coeffs = extract_affine(v)
if (isinstance(const, Tensor) and all(
isinstance(coeff, Tensor)
for coeff, _ in coeffs.values())):
affine[k] = const, coeffs
else:
remaining_subs += (k, v),
if not affine:
return reflect(Subs, self, remaining_subs)
# Align integer dimensions.
old_int_inputs = OrderedDict(
(k, v) for k, v in self.inputs.items() if v.dtype != 'real')
tensors = [
Tensor(self.info_vec, old_int_inputs),
Tensor(self.precision, old_int_inputs)
]
for const, coeffs in affine.values():
tensors.append(const)
tensors.extend(coeff for coeff, _ in coeffs.values())
new_int_inputs, tensors = align_tensors(*tensors, expand=True)
tensors = (Tensor(x, new_int_inputs) for x in tensors)
old_info_vec = next(tensors).data
old_precision = next(tensors).data
for old_k, (const, coeffs) in affine.items():
const = next(tensors)
for new_k, (coeff, eqn) in coeffs.items():
coeff = next(tensors)
coeffs[new_k] = coeff, eqn
affine[old_k] = const, coeffs
batch_shape = old_info_vec.data.shape[:-1]
# Align real dimensions.
old_real_inputs = OrderedDict(
(k, v) for k, v in self.inputs.items() if v.dtype == 'real')
new_real_inputs = old_real_inputs.copy()
for old_k, (const, coeffs) in affine.items():
del new_real_inputs[old_k]
for new_k, (coeff, eqn) in coeffs.items():
new_shape = coeff.shape[:len(eqn.split('->')[0].split(',')[1])]
new_real_inputs[new_k] = reals(*new_shape)
old_offsets, old_dim = _compute_offsets(old_real_inputs)
new_offsets, new_dim = _compute_offsets(new_real_inputs)
new_inputs = new_int_inputs.copy()
new_inputs.update(new_real_inputs)
# Construct a blockwise affine representation of the substitution.
subs_vector = BlockVector(batch_shape + (old_dim, ))
subs_matrix = BlockMatrix(batch_shape + (new_dim, old_dim))
for old_k, old_offset in old_offsets.items():
old_size = old_real_inputs[old_k].num_elements
old_slice = slice(old_offset, old_offset + old_size)
if old_k in new_real_inputs:
new_offset = new_offsets[old_k]
new_slice = slice(new_offset, new_offset + old_size)
subs_matrix[..., new_slice, old_slice] = \
torch.eye(old_size).expand(batch_shape + (-1, -1))
continue
const, coeffs = affine[old_k]
old_shape = old_real_inputs[old_k].shape
assert const.data.shape == batch_shape + old_shape
subs_vector[..., old_slice] = const.data.reshape(batch_shape +
(old_size, ))
for new_k, new_offset in new_offsets.items():
if new_k in coeffs:
coeff, eqn = coeffs[new_k]
new_size = new_real_inputs[new_k].num_elements
new_slice = slice(new_offset, new_offset + new_size)
assert coeff.shape == new_real_inputs[
new_k].shape + old_shape
subs_matrix[..., new_slice, old_slice] = \
coeff.data.reshape(batch_shape + (new_size, old_size))
subs_vector = subs_vector.as_tensor()
subs_matrix = subs_matrix.as_tensor()
subs_matrix_t = subs_matrix.transpose(-1, -2)
# Construct the new funsor. Suppose the old Gaussian funsor g has density
# g(x) = < x | i - 1/2 P x>
# Now define a new funsor f by substituting x = A y + B:
# f(y) = g(A y + B)
# = < A y + B | i - 1/2 P (A y + B) >
# = < y | At (i - P B) - 1/2 At P A y > + < B | i - 1/2 P B >
# =: < y | i' - 1/2 P' y > + C
# where P' = At P A and i' = At (i - P B) parametrize a new Gaussian
# and C = < B | i - 1/2 P B > parametrize a new Tensor.
precision = subs_matrix @ old_precision @ subs_matrix_t
info_vec = _mv(subs_matrix,
old_info_vec - _mv(old_precision, subs_vector))
const = _vv(subs_vector,
old_info_vec - 0.5 * _mv(old_precision, subs_vector))
result = Gaussian(info_vec, precision, new_inputs) + Tensor(
const, new_int_inputs)
return Subs(result, remaining_subs) if remaining_subs else result
def test_mvn_defaults():
loc = Variable('loc', reals(3))
scale_tril = Variable('scale', reals(3, 3))
value = Variable('value', reals(3))
assert dist.MultivariateNormal(loc, scale_tril) is dist.MultivariateNormal(
loc, scale_tril, value)
def test_eager_subs_variable():
v = Variable('v', reals(3))
point = Tensor(torch.randn(3))
d = Delta('foo', v)
assert d(v=point) is Delta('foo', point)
def unscaled_sample(self, sampled_vars, sample_inputs):
assert self.output == reals()
sampled_vars = sampled_vars.intersection(self.inputs)
if not sampled_vars:
return self
# Partition inputs into sample_inputs + batch_inputs + event_inputs.
sample_inputs = OrderedDict(
(k, d) for k, d in sample_inputs.items() if k not in self.inputs)
sample_shape = tuple(int(d.dtype) for d in sample_inputs.values())
batch_inputs = OrderedDict(
(k, d) for k, d in self.inputs.items() if k not in sampled_vars)
event_inputs = OrderedDict(
(k, d) for k, d in self.inputs.items() if k in sampled_vars)
be_inputs = batch_inputs.copy()
be_inputs.update(event_inputs)
sb_inputs = sample_inputs.copy()
sb_inputs.update(batch_inputs)
# Sample all variables in a single Categorical call.
logits = align_tensor(be_inputs, self)
batch_shape = logits.shape[:len(batch_inputs)]
flat_logits = logits.reshape(batch_shape + (-1, ))
sample_shape = tuple(d.dtype for d in sample_inputs.values())
flat_sample = torch.distributions.Categorical(
logits=flat_logits).sample(sample_shape)
assert flat_sample.shape == sample_shape + batch_shape
results = []
mod_sample = flat_sample
for name, domain in reversed(list(event_inputs.items())):
size = domain.dtype
point = Tensor(mod_sample % size, sb_inputs, size)
mod_sample = mod_sample / size
results.append(Delta(name, point))
# Account for the log normalizer factor.
# Derivation: Let f be a nonnormalized distribution (a funsor), and
# consider operations in linear space (source code is in log space).
# Let x0 ~ f/|f| be a monte carlo sample from a normalized f/|f|.
# f(x0) / |f| # dice numerator
# Let g = delta(x=x0) |f| -----------------
# detach(f(x0)/|f|) # dice denominator
# |detach(f)| f(x0)
# = delta(x=x0) ----------------- be a dice approximation of f.
# detach(f(x0))
# Then g is an unbiased estimator of f in value and all derivatives.
# In the special case f = detach(f), we can simplify to
# g = delta(x=x0) |f|.
if flat_logits.requires_grad:
# Apply a dice factor to preserve differentiability.
index = [
torch.arange(n).reshape((n, ) + (1, ) *
(flat_logits.dim() - i - 2))
for i, n in enumerate(flat_logits.shape[:-1])
]
index.append(flat_sample)
log_prob = flat_logits[index]
assert log_prob.shape == flat_sample.shape
results.append(
Tensor(
flat_logits.detach().logsumexp(-1) +
(log_prob - log_prob.detach()), sb_inputs))
else:
# This is the special case f = detach(f).
results.append(Tensor(flat_logits.logsumexp(-1), batch_inputs))
return reduce(ops.add, results)
def test_eager_subs_ground(log_density):
point1 = Tensor(torch.randn(3))
point2 = Tensor(torch.randn(3))
d = Delta('foo', point1, log_density)
check_funsor(d(foo=point1), {}, reals(), torch.tensor(float(log_density)))
check_funsor(d(foo=point2), {}, reals(), torch.tensor(float('-inf')))
def test_to_data_error():
with pytest.raises(ValueError):
to_data(Variable('x', reals()))
with pytest.raises(ValueError):
to_data(Variable('y', bint(12)))
def test_transform_log(shape):
point = Tensor(torch.randn(shape))
x = Variable('x', reals(*shape))
actual = Delta('y', point)(y=ops.log(x))
expected = Delta('x', point.exp(), -point.sum())
assert_close(actual, expected)
x = Variable('x', reals())
assert isinstance(x, Variable)
y = Variable('y', reals())
assert isinstance(y, Variable)
result = eval(expr)
assert isinstance(result, expected_type)
assert result.is_affine
SUBS_TESTS = [
("(t * x)(i=1)", Contraction, {
"j": bint(3),
"x": reals()
}),
("(t * x)(i=1, x=y)", Contraction, {
"j": bint(3),
"y": reals()
}),
("(t * x + n)(x=y)", Contraction, {
"y": reals(),
"i": bint(2),
"j": bint(3)
}),
("(x + y)(y=z)", Contraction, {
"x": reals(),
"z": reals()
}),
("(-x)(x=y+z)", Contraction, {
def test_binomial_density(batch_shape, eager):
batch_dims = ('i', 'j', 'k')[:len(batch_shape)]
inputs = OrderedDict((k, bint(v)) for k, v in zip(batch_dims, batch_shape))
max_count = 10
@funsor.torch.function(reals(), reals(), reals(), reals())
def binomial(total_count, probs, value):
return torch.distributions.Binomial(total_count, probs).log_prob(value)
check_funsor(binomial, {
'total_count': reals(),
'probs': reals(),
'value': reals()
}, reals())
value_data = random_tensor(inputs, bint(max_count)).data.float()
total_count_data = value_data + random_tensor(
inputs, bint(max_count)).data.float()
value = Tensor(value_data, inputs)
total_count = Tensor(total_count_data, inputs)
probs = Tensor(torch.rand(batch_shape), inputs)
expected = binomial(total_count, probs, value)
check_funsor(expected, inputs, reals())
m = Variable('value', reals())
actual = dist.Binomial(total_count, probs, value) if eager else \
dist.Binomial(total_count, probs, m)(value=value)
check_funsor(actual, inputs, reals())
assert_close(actual, expected)
factors1 = partial_sum_product(sum_op, prod_op, factors, vars1, plates)
factors2 = partial_sum_product(sum_op, prod_op, factors1, vars2, plates)
actual = reduce(prod_op, factors2)
expected = sum_product(sum_op, prod_op, factors, vars1 | vars2, plates)
assert_close(actual, expected)
@pytest.mark.parametrize('num_steps', [None] + list(range(1, 13)))
@pytest.mark.parametrize('sum_op,prod_op,state_domain', [
(ops.add, ops.mul, bint(2)),
(ops.add, ops.mul, bint(3)),
(ops.logaddexp, ops.add, bint(2)),
(ops.logaddexp, ops.add, bint(3)),
(ops.logaddexp, ops.add, reals()),
(ops.logaddexp, ops.add, reals(2)),
],
ids=str)
@pytest.mark.parametrize('batch_inputs', [
{},
{
"foo": bint(5)
},
{
"foo": bint(2),
"bar": bint(4)
},
],
ids=lambda d: ",".join(d.keys()))
@pytest.mark.parametrize('impl', [
def test_categorical_defaults():
probs = Variable('probs', reals(3))
value = Variable('value', bint(3))
assert dist.Categorical(probs) is dist.Categorical(probs, value)
@pytest.mark.parametrize('int_inputs', [
{},
{
'i': bint(2)
},
{
'i': bint(2),
'j': bint(3)
},
],
ids=id_from_inputs)
@pytest.mark.parametrize('real_inputs', [
{
'x': reals()
},
{
'x': reals(4)
},
{
'x': reals(2, 3)
},
{
'x': reals(),
'y': reals()
},
{
'x': reals(2),
'y': reals(3)
},
def test_affine_subs():
# This was recorded from test_pyro_convert.
x = Subs(
Gaussian(
torch.tensor([
1.3027106523513794, 1.4167094230651855, -0.9750942587852478,
0.5321089029312134, -0.9039931297302246
],
dtype=torch.float32), # noqa
torch.tensor([[
1.0199567079544067, 0.9840421676635742, -0.473368763923645,
0.34206756949424744, -0.7562517523765564
],
[
0.9840421676635742, 1.511502742767334,
-1.7593903541564941, 0.6647964119911194,
-0.5119513273239136
],
[
-0.4733688533306122, -1.7593903541564941,
3.2386727333068848, -0.9345928430557251,
-0.1534711718559265
],
[
0.34206756949424744, 0.6647964119911194,
-0.9345928430557251, 0.3141004145145416,
-0.12399007380008698
],
[
-0.7562517523765564, -0.5119513273239136,
-0.1534711718559265, -0.12399007380008698,
0.6450173854827881
]],
dtype=torch.float32), # noqa
(
(
'state_1_b6',
reals(3, ),
),
(
'obs_b2',
reals(2, ),
),
)),
(
(
'obs_b2',
Contraction(
ops.nullop,
ops.add,
frozenset(),
(
Variable('bias_b5', reals(2, )),
Tensor(
torch.tensor(
[-2.1787893772125244, 0.5684312582015991],
dtype=torch.float32), # noqa
(),
'real'),
)),
), ))
assert isinstance(x, (Gaussian, Contraction)), x.pretty()
