def call(self, inputs):
net = self.layer(inputs)
logits, loc, unconstrained_scale = tf.split(net, 3, axis=-1)
scale = tf.nn.softplus(unconstrained_scale) + tf.keras.backend.epsilon()
return ed.MixtureSameFamily(
mixture_distribution=ed.Categorical(logits=logits).distribution,
components_distribution=ed.Logistic(loc=loc, scale=scale).distribution)
def model_mixture_adaptive2(X, ls=1., n_mix=2, ridge_factor=1e-3):
"""Alternative representation using Mixture family.
Args:
X: (np.ndarray of float32) input training features.
with dimension (N, D).
ls: (float32) length scale parameter.
n_mix: (int8) Number of mixture components.
ridge_factor: (float32) ridge factor to stabilize Cholesky decomposition.
Returns:
(tf.Tensors of float32) model parameters.
"""
N = X.shape[0]
K_mat = rbf(X, ls=ls, ridge_factor=ridge_factor)
gp_weight = ed.Independent(distribution=tfd.MultivariateNormalTriL(
loc=tf.zeros(shape=[n_mix, N]),
scale_tril=replicate_along_zero_axis(tf.cholesky(K_mat), n_mix),
),
reinterpreted_batch_ndims=1,
name="gp_w")
mix_member = ed.Multinomial(total_count=[1.],
logits=tf.transpose(gp_weight),
name="mix_prob")
gp_comp = ed.Independent(distribution=tfd.MultivariateNormalTriL(
loc=tf.zeros(shape=[n_mix, N]),
scale_tril=replicate_along_zero_axis(tf.cholesky(K_mat), n_mix),
),
reinterpreted_batch_ndims=1,
name="gp_f")
sigma = ed.Normal(loc=tf.ones(n_mix) * -5.,
scale=tf.ones(n_mix) * 1.,
name='sigma')
y = ed.MixtureSameFamily(
components_distribution=tfd.MultivariateNormalDiag(
loc=gp_comp, scale_identity_multiplier=tf.exp(sigma)),
mixture_distribution=tfd.Categorical(logits=tf.transpose(gp_weight)),
name="y")
return gp_weight, mix_member, gp_comp, sigma, y
def model_mixture(X, ls=1., n_mix=2, ridge_factor=1e-3):
"""Defines the Gaussian Process Model.
Args:
X: (np.ndarray of float32) input training features.
with dimension (N, D).
ls: (float32) length scale parameter.
n_mix: (int8) Number of mixture components.
ridge_factor: (float32) ridge factor to stabilize Cholesky decomposition.
Returns:
(tf.Tensors of float32) model parameters.
"""
N = X.shape[0]
K_mat = rbf(X, ls=ls, ridge_factor=ridge_factor)
mix_prob = ed.Dirichlet(concentration=tf.ones(n_mix, dtype=tf.float32) /
n_mix,
name='mix_prob')
gp_f = ed.Independent(distribution=tfd.MultivariateNormalTriL(
loc=tf.zeros(shape=[n_mix, N]),
scale_tril=replicate_along_zero_axis(tf.cholesky(K_mat), n_mix),
),
reinterpreted_batch_ndims=1,
name="gp_f")
sigma = ed.Normal(loc=tf.ones(n_mix) * -5.,
scale=tf.ones(n_mix) * 1.,
name='sigma')
y = ed.MixtureSameFamily(
components_distribution=tfd.MultivariateNormalDiag(
loc=gp_f, scale_identity_multiplier=tf.exp(sigma)),
mixture_distribution=tfd.Categorical(probs=mix_prob),
name="y")
return mix_prob, gp_f, sigma, y