def to_dist_fn(h):
mu, ln_sigma = h
xp = cuda.get_array_module(*h)
scale = F.softplus(ln_sigma)
return distributions.HyperbolicWrapped(
distributions.Independent(D.Normal(
loc=xp.zeros(shape=scale.shape, dtype=scale.dtype),
scale=scale)),
functions.pseudo_polar_projection(mu))
def to_dist_fn(h):
loc, scale = h
xp = cuda.get_array_module(loc, scale)
return distributions.HyperbolicWrapped(
distributions.Independent(
D.Normal(loc=xp.zeros(shape=scale.shape,
dtype=scale.dtype),
scale=scale)),
functions.pseudo_polar_projection(loc))
def to_dist_fn(h):
xp = cuda.get_array_module(h)
scale = F.softplus(h[..., n_latent:])
shape = scale.shape[:-1] + (n_latent, )
return distributions.HyperbolicWrapped(
distributions.Independent(
D.Normal(loc=xp.zeros(shape=shape, dtype=scale.dtype),
scale=scale)),
functions.pseudo_polar_projection(h[..., :n_latent]))
def to_dist_fn(h):
if ndim == 1:
dim_h = h.shape[-1]
loc = h[..., :(dim_h // 2)]
base_sigma = h[..., (dim_h // 2):]
elif ndim == 3:
nb_channel = h.shape[-3]
loc = h[..., :(nb_channel // 2), :, :]
base_sigma = h[..., (nb_channel // 2):, :, :]
else:
raise NotImplementedError
base_sigma += xp_functions._softplus_inverse(1.0)
return distributions.Independent(D.Normal(
loc=loc,
scale=F.softplus(
functions.clamp(base_sigma,
xp_functions._softplus_inverse(0.001)))),
reinterpreted_batch_ndims=ndim)
def to_dist_fn(h):
mu, ln_sigma = h
return distributions.Independent(D.Normal(
loc=mu, scale=F.softplus(ln_sigma)))
def to_dist_fn(h):
loc, scale = h
return distributions.Independent(D.Normal(loc, scale),
reinterpreted_batch_ndims=1)
def to_dist_fn(h):
return distributions.Independent(D.Bernoulli(logit=h),
reinterpreted_batch_ndims=ndim)
def to_dist_fn(h):
return distributions.Independent(
D.Normal(loc=h[..., :n_latent],
scale=F.softplus(h[..., n_latent:])))