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