jit_decoder = pyro.ops.jit.trace(decoder)
states = decoder(transition, means, data)
jit_states = jit_decoder(transition, means, data)
assert len(states) == len(jit_states)
for state, jit_state in zip(states, jit_states):
assert state.shape == jit_state.shape
if temperature == 0:
assert_equal(state, jit_state)
@pytest.mark.parametrize(
"x,y",
[
(
CondIndepStackFrame("a", -1, torch.tensor(2000), 2),
CondIndepStackFrame("a", -1, 2000, 2),
),
(
CondIndepStackFrame("a", -1, 1, 2),
CondIndepStackFrame("a", -1, torch.tensor(1), 2),
),
],
)
def test_cond_indep_equality(x, y):
assert x == y
assert not x != y
assert hash(x) == hash(y)
def test_jit_arange_workaround():
def _pyro_param(self, msg):
frame = CondIndepStackFrame(self.name, None, self.size, self._counter)
msg["cond_indep_stack"] = (frame, ) + msg["cond_indep_stack"]
msg["scale"] = msg["scale"] * self._scale
def _pyro_param(self, msg):
frame = CondIndepStackFrame(self.name, self.dim, self.size, 0)
msg["cond_indep_stack"] = (frame, ) + msg["cond_indep_stack"]
def frame(dim, size):
return CondIndepStackFrame(name="plate_{}".format(size), dim=dim, size=size, counter=0)