def _get_transform(self):
flows = []
for i in range(self.num_flows):
arn_params = param('{}_arn__{}'.format(self.prefix, i), None)
if i > 0:
flows.append(
PermuteTransform(np.arange(self.latent_size)[::-1]))
flows.append(
InverseAutoregressiveTransform(self.arns[i], arn_params))
return ComposeTransform(flows)
def test_iaf():
# test for substitute logic for exposed methods `sample_posterior` and `get_transforms`
N, dim = 3000, 3
data = random.normal(random.PRNGKey(0), (N, dim))
true_coefs = jnp.arange(1.0, dim + 1.0)
logits = jnp.sum(true_coefs * data, axis=-1)
labels = dist.Bernoulli(logits=logits).sample(random.PRNGKey(1))
def model(data, labels):
coefs = numpyro.sample("coefs",
dist.Normal(jnp.zeros(dim), jnp.ones(dim)))
offset = numpyro.sample("offset", dist.Uniform(-1, 1))
logits = offset + jnp.sum(coefs * data, axis=-1)
return numpyro.sample("obs", dist.Bernoulli(logits=logits), obs=labels)
adam = optim.Adam(0.01)
rng_key_init = random.PRNGKey(1)
guide = AutoIAFNormal(model)
svi = SVI(model, guide, adam, Trace_ELBO())
svi_state = svi.init(rng_key_init, data, labels)
params = svi.get_params(svi_state)
x = random.normal(random.PRNGKey(0), (dim + 1, ))
rng_key = random.PRNGKey(1)
actual_sample = guide.sample_posterior(rng_key, params)
actual_output = guide._unpack_latent(guide.get_transform(params)(x))
flows = []
for i in range(guide.num_flows):
if i > 0:
flows.append(transforms.PermuteTransform(
jnp.arange(dim + 1)[::-1]))
arn_init, arn_apply = AutoregressiveNN(
dim + 1,
[dim + 1, dim + 1],
permutation=jnp.arange(dim + 1),
skip_connections=guide._skip_connections,
nonlinearity=guide._nonlinearity,
)
arn = partial(arn_apply, params["auto_arn__{}$params".format(i)])
flows.append(InverseAutoregressiveTransform(arn))
flows.append(guide._unpack_latent)
transform = transforms.ComposeTransform(flows)
_, rng_key_sample = random.split(rng_key)
expected_sample = transform(
dist.Normal(jnp.zeros(dim + 1), 1).sample(rng_key_sample))
expected_output = transform(x)
assert_allclose(actual_sample["coefs"], expected_sample["coefs"])
assert_allclose(
actual_sample["offset"],
transforms.biject_to(constraints.interval(-1, 1))(
expected_sample["offset"]),
)
check_eq(actual_output, expected_output)
def test_iaf():
# test for substitute logic for exposed methods `sample_posterior` and `get_transforms`
N, dim = 3000, 3
data = random.normal(random.PRNGKey(0), (N, dim))
true_coefs = np.arange(1., dim + 1.)
logits = np.sum(true_coefs * data, axis=-1)
labels = dist.Bernoulli(logits=logits).sample(random.PRNGKey(1))
def model(data, labels):
coefs = numpyro.sample('coefs', dist.Normal(np.zeros(dim),
np.ones(dim)))
offset = numpyro.sample('offset', dist.Uniform(-1, 1))
logits = offset + np.sum(coefs * data, axis=-1)
return numpyro.sample('obs', dist.Bernoulli(logits=logits), obs=labels)
adam = optim.Adam(0.01)
rng_init = random.PRNGKey(1)
guide = AutoIAFNormal(model)
svi = SVI(model, guide, elbo, adam)
svi_state = svi.init(rng_init,
model_args=(data, labels),
guide_args=(data, labels))
params = svi.get_params(svi_state)
x = random.normal(random.PRNGKey(0), (dim + 1, ))
rng = random.PRNGKey(1)
actual_sample = guide.sample_posterior(rng, params)
actual_output = guide.get_transform(params)(x)
flows = []
for i in range(guide.num_flows):
if i > 0:
flows.append(constraints.PermuteTransform(
np.arange(dim + 1)[::-1]))
arn_init, arn_apply = AutoregressiveNN(
dim + 1, [dim + 1, dim + 1],
permutation=np.arange(dim + 1),
skip_connections=guide._skip_connections,
nonlinearity=guide._nonlinearity)
arn = partial(arn_apply, params['auto_arn__{}$params'.format(i)])
flows.append(InverseAutoregressiveTransform(arn))
transform = constraints.ComposeTransform(flows)
rng_seed, rng_sample = random.split(rng)
expected_sample = guide.unpack_latent(
transform(dist.Normal(np.zeros(dim + 1), 1).sample(rng_sample)))
expected_output = transform(x)
assert_allclose(actual_sample['coefs'], expected_sample['coefs'])
assert_allclose(
actual_sample['offset'],
constraints.biject_to(constraints.interval(-1, 1))(
expected_sample['offset']))
assert_allclose(actual_output, expected_output)
def _get_transform(self):
if self.latent_size == 1:
raise ValueError('latent dim = 1. Consider using AutoDiagonalNormal instead')
hidden_dims = [self.latent_size, self.latent_size] if self._hidden_dims is None else self._hidden_dims
flows = []
for i in range(self.num_flows):
if i > 0:
flows.append(PermuteTransform(np.arange(self.latent_size)[::-1]))
arn = AutoregressiveNN(self.latent_size, hidden_dims,
permutation=np.arange(self.latent_size),
skip_connections=self._skip_connections,
nonlinearity=self._nonlinearity)
arnn = numpyro.module('{}_arn__{}'.format(self.prefix, i), arn, (self.latent_size,))
flows.append(InverseAutoregressiveTransform(arnn))
return ComposeTransform(flows)
def _get_posterior(self):
if self.latent_dim == 1:
raise ValueError('latent dim = 1. Consider using AutoDiagonalNormal instead')
hidden_dims = [self.latent_dim, self.latent_dim] if self._hidden_dims is None else self._hidden_dims
flows = []
for i in range(self.num_flows):
if i > 0:
flows.append(PermuteTransform(jnp.arange(self.latent_dim)[::-1]))
arn = AutoregressiveNN(self.latent_dim, hidden_dims,
permutation=jnp.arange(self.latent_dim),
skip_connections=self._skip_connections,
nonlinearity=self._nonlinearity)
arnn = numpyro.module('{}_arn__{}'.format(self.prefix, i), arn, (self.latent_dim,))
flows.append(InverseAutoregressiveTransform(arnn))
return dist.TransformedDistribution(self.get_base_dist(), flows)
def _make_iaf(input_dim, hidden_dims, rng):
arn_init, arn = AutoregressiveNN(input_dim, hidden_dims, param_dims=[1, 1])
_, init_params = arn_init(rng, (input_dim, ))
return InverseAutoregressiveTransform(partial(arn, init_params))
