def p_ugz(self):
if hasattr(self, "_pugz"):
return self._pugz
# Keep this simple!
self._pugz = nux.ParametrizedGaussianPrior(
network_kwargs=self.network_kwargs,
create_network=self.create_network)
return self._pugz
def q_ugx(self):
if hasattr(self, "_qugx"):
return self._qugx
# Keep this simple, but a bit more complicated than p(u|z).
self._qugx = nux.sequential(
nux.reverse_flow(
nux.LogisticMixtureLogit(n_components=8,
with_affine_coupling=False,
coupling=False)),
nux.ParametrizedGaussianPrior(network_kwargs=self.network_kwargs,
create_network=self.create_network))
return self._qugx
def default_flow(self):
def block():
return nux.sequential(nux.RationalQuadraticSpline(K=8,
network_kwargs=self.network_kwargs,
create_network=self.create_network,
use_condition=True,
coupling=True,
condition_method="nin"),
nux.OneByOneConv())
f = nux.repeat(block, n_repeats=3)
return nux.sequential(f,
nux.ParametrizedGaussianPrior(network_kwargs=self.network_kwargs,
create_network=self.create_network))
def default_flow(self):
def block():
return nux.sequential(nux.RationalQuadraticSpline(K=8,
network_kwargs=self.network_kwargs,
create_network=self.create_network,
use_condition=True,
coupling=True,
condition_method="concat"),
nux.AffineLDU(safe_diag=True))
f = nux.repeat(block, n_repeats=3)
return nux.sequential(nux.reverse_flow(f),
nux.ParametrizedGaussianPrior(network_kwargs=self.network_kwargs,
create_network=self.create_network))