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):
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):
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]))