Ops: keras_proba.ops

Tensor operations on the outputs and parameters of keras_proba layers.

mean_field_logits

mean_field_logits(logits, covariance_matrix=None, mean_field_factor=1.0)

Adjusts the model logits so their softmax approximates the posterior mean [1].

[1]: Zhiyun Lu, Eugene Ie, Fei Sha. Uncertainty Estimation with Infinitesimal Jackknife. arXiv preprint arXiv:2006.07584, 2020. https://arxiv.org/abs/2006.07584

Parameters:

Name	Type	Description	Default
logits		Float tensor of shape (batch_size, num_classes).	required
covariance_matrix		Covariance matrix of shape (batch_size, batch_size). If None, it is assumed to be the identity.	None
mean_field_factor		Scale factor of the mean-field approximation, which controls the influence of the posterior variance. If covariance_matrix=None, it acts as a temperature for temperature scaling. If None or negative, logits are returned unchanged.	1.0

Returns:

Type	Description
	Adjusted logits of shape (batch_size, num_classes).

split_gaussian_params

split_gaussian_params(params)

Splits a GaussianHead output into (mean, scale).

Parameters:

Name	Type	Description	Default
params		Tensor of shape (..., 2 * units).	required

Returns:

Type	Description
	Tuple of mean and scale tensors, each of shape (..., units).

split_mixture_params

split_mixture_params(params, units=1)

Splits a MixtureGaussianHead output into (logits, means, scales).

Parameters:

Name	Type	Description	Default
params		Tensor of shape (..., K * (1 + 2 * units)).	required
units		Number of regression targets D.	1

Returns:

Type	Description
	Tuple of logits of shape (..., K), and means and scales of shape
	(..., K, units).

mixture_mean_and_stddev

mixture_mean_and_stddev(params, units=1)

Predictive mean and standard deviation of a MixtureGaussianHead output.

The variance is the law of total variance, sum_k pi_k (scale_k^2 + mean_k^2) - mean^2, so it includes the spread between components.

Parameters:

Name	Type	Description	Default
params		Tensor of shape (..., K * (1 + 2 * units)).	required
units		Number of regression targets.	1

Returns:

Type	Description
	Tuple of mean and stddev, each of shape (..., units).

normal_kl_divergence

normal_kl_divergence(loc, scale, prior_scale)

KL(N(loc, scale^2) || N(0, prior_scale^2)), summed over all elements.