Skip to content

Module fl_server_ai.uncertainty

View Source
# SPDX-FileCopyrightText: 2024 Benedikt Franke <benedikt.franke@dlr.de>
# SPDX-FileCopyrightText: 2024 Florian Heinrich <florian.heinrich@dlr.de>
#
# SPDX-License-Identifier: Apache-2.0

from .base import UncertaintyBase
from .ensemble import Ensemble
from .mc_dropout import MCDropout
from .method import get_uncertainty_class
from .none import NoneUncertainty
from .swag import SWAG


__all__ = [
    "get_uncertainty_class",
    "Ensemble",
    "MCDropout",
    "NoneUncertainty",
    "SWAG",
    "UncertaintyBase",
]

Sub-modules

Functions

get_uncertainty_class

def get_uncertainty_class(
    value: fl_server_core.models.model.Model | fl_server_core.models.training.Training | fl_server_core.models.training.UncertaintyMethod
) -> Type[fl_server_ai.uncertainty.base.UncertaintyBase]

Get uncertainty class associated with a given Model, Training, or UncertaintyMethod object.

Parameters:

Name Type Description Default
value Model Training UncertaintyMethod

Returns:

Type Description
Type[UncertaintyBase] The uncertainty class associated with the given object.

Raises:

Type Description
ValueError If the given object is not a Model, Training, or UncertaintyMethod,
or if the uncertainty method associated with the object is unknown.
View Source
def get_uncertainty_class(value: Model | Training | UncertaintyMethod) -> Type[UncertaintyBase]:
    """
    Get uncertainty class associated with a given Model, Training, or UncertaintyMethod object.

    Args:
        value (Model | Training | UncertaintyMethod): The object to retrieve the uncertainty class for.

    Returns:
        Type[UncertaintyBase]: The uncertainty class associated with the given object.

    Raises:
        ValueError: If the given object is not a Model, Training, or UncertaintyMethod,
                    or if the uncertainty method associated with the object is unknown.
    """
    if isinstance(value, UncertaintyMethod):
        method = value
    elif isinstance(value, Training):
        method = value.uncertainty_method
    elif isinstance(value, Model):
        uncertainty_method = Training.objects.filter(model=value) \
                .values("uncertainty_method") \
                .first()["uncertainty_method"]
        method = uncertainty_method
    else:
        raise ValueError(f"Unknown type: {type(value)}")

    match method:
        case UncertaintyMethod.ENSEMBLE: return Ensemble
        case UncertaintyMethod.MC_DROPOUT: return MCDropout
        case UncertaintyMethod.NONE: return NoneUncertainty
        case UncertaintyMethod.SWAG: return SWAG
        case _: raise ValueError(f"Unknown uncertainty method: {method}")

Classes

Ensemble

class Ensemble(
    /,
    *args,
    **kwargs
)

Ensemble uncertainty estimation.

View Source
class Ensemble(UncertaintyBase):
    """
    Ensemble uncertainty estimation.
    """

    @classmethod
    def prediction(cls, input: torch.Tensor, model: MeanModel) -> Tuple[torch.Tensor, Dict[str, Any]]:
        output_list = []
        for m in model.models.all():
            net = m.get_torch_model()
            output = net(input).detach()
            output_list.append(output)
        outputs = torch.stack(output_list, dim=0)  # (N, batch_size, n_classes)  # N = number of models

        inference = outputs.mean(dim=0)
        uncertainty = cls.interpret(outputs)
        return inference, uncertainty

Ancestors (in MRO)

  • fl_server_ai.uncertainty.UncertaintyBase
  • abc.ABC

Static methods

expected_entropy

def expected_entropy(
    predictions: torch.Tensor
) -> torch.Tensor

Calculate the mean entropy of the predictive distribution across the MC samples.

Parameters:

Name Type Description Default
predictions torch.Tensor predictions of shape (n_mc x batch_size x n_classes) None

Returns:

Type Description
torch.Tensor mean entropy of the predictive distribution
View Source
    @classmethod
    def expected_entropy(cls, predictions: torch.Tensor) -> torch.Tensor:
        """
        Calculate the mean entropy of the predictive distribution across the MC samples.

        Args:
            predictions (torch.Tensor): predictions of shape (n_mc x batch_size x n_classes)

        Returns:
            torch.Tensor: mean entropy of the predictive distribution
        """
        return torch.distributions.Categorical(probs=predictions).entropy().mean(dim=0)

get_options

def get_options(
    obj: fl_server_core.models.model.Model | fl_server_core.models.training.Training
) -> Dict[str, Any]

Get uncertainty options from training options.

Parameters:

Name Type Description Default
obj Model Training The Model or Training object to retrieve options for.

Returns:

Type Description
Dict[str, Any] Uncertainty options.

Raises:

Type Description
TypeError If the given object is not a Model or Training.
View Source
    @classmethod
    def get_options(cls, obj: Model | Training) -> Dict[str, Any]:
        """
        Get uncertainty options from training options.

        Args:
            obj (Model | Training): The Model or Training object to retrieve options for.

        Returns:
            Dict[str, Any]: Uncertainty options.

        Raises:
            TypeError: If the given object is not a Model or Training.
        """
        if isinstance(obj, Model):
            return Training.objects.filter(model=obj) \
                .values("options") \
                .first()["options"] \
                .get("uncertainty", {})
        if isinstance(obj, Training):
            return obj.options.get("uncertainty", {})
        raise TypeError(f"Expected Model or Training, got {type(obj)}")

interpret

def interpret(
    outputs: torch.Tensor
) -> Dict[str, Any]

Interpret the different network (model) outputs and calculate the uncertainty.

Parameters:

Name Type Description Default
outputs torch.Tensor multiple network (model) outputs (N, batch_size, n_classes) None
View Source
    @classmethod
    def interpret(cls, outputs: torch.Tensor) -> Dict[str, Any]:
        """
        Interpret the different network (model) outputs and calculate the uncertainty.

        Args:
            outputs (torch.Tensor): multiple network (model) outputs (N, batch_size, n_classes)

        Return:
            Tuple[torch.Tensor, Dict[str, Any]]: inference and uncertainty
        """
        variance = outputs.var(dim=0)
        std = outputs.std(dim=0)
        if not (torch.all(outputs <= 1.) and torch.all(outputs >= 0.)):
            return dict(variance=variance, std=std)

        predictive_entropy = cls.predictive_entropy(outputs)
        expected_entropy = cls.expected_entropy(outputs)
        mutual_info = predictive_entropy - expected_entropy  # see cls.mutual_information
        return dict(
            variance=variance,
            std=std,
            predictive_entropy=predictive_entropy,
            expected_entropy=expected_entropy,
            mutual_info=mutual_info,
        )

mutual_information

def mutual_information(
    predictions: torch.Tensor
) -> torch.Tensor

Calculate the BALD (Bayesian Active Learning by Disagreement) of a model;

the difference between the mean of the entropy and the entropy of the mean of the predicted distribution on the predictions. This method is also sometimes referred to as the mutual information (MI).

Parameters:

Name Type Description Default
predictions torch.Tensor predictions of shape (n_mc x batch_size x n_classes) None

Returns:

Type Description
torch.Tensor difference between the mean of the entropy and the entropy of the mean
of the predicted distribution
View Source
    @classmethod
    def mutual_information(cls, predictions: torch.Tensor) -> torch.Tensor:
        """
        Calculate the BALD (Bayesian Active Learning by Disagreement) of a model;
        the difference between the mean of the entropy and the entropy of the mean
        of the predicted distribution on the predictions.
        This method is also sometimes referred to as the mutual information (MI).

        Args:
            predictions (torch.Tensor): predictions of shape (n_mc x batch_size x n_classes)

        Returns:
            torch.Tensor: difference between the mean of the entropy and the entropy of the mean
                    of the predicted distribution
        """
        return cls.predictive_entropy(predictions) - cls.expected_entropy(predictions)

prediction

def prediction(
    input: torch.Tensor,
    model: fl_server_core.models.model.MeanModel
) -> Tuple[torch.Tensor, Dict[str, Any]]

Make a prediction using the given input and model.

Parameters:

Name Type Description Default
input torch.Tensor The input to make a prediction for. None
model Model The model to use for making the prediction. None

Returns:

Type Description
Tuple[torch.Tensor, Dict[str, Any]] The prediction and any associated uncertainty.
View Source
    @classmethod
    def prediction(cls, input: torch.Tensor, model: MeanModel) -> Tuple[torch.Tensor, Dict[str, Any]]:
        output_list = []
        for m in model.models.all():
            net = m.get_torch_model()
            output = net(input).detach()
            output_list.append(output)
        outputs = torch.stack(output_list, dim=0)  # (N, batch_size, n_classes)  # N = number of models

        inference = outputs.mean(dim=0)
        uncertainty = cls.interpret(outputs)
        return inference, uncertainty

predictive_entropy

def predictive_entropy(
    predictions: torch.Tensor
) -> torch.Tensor

Calculate the entropy of the mean predictive distribution across the MC samples.

Parameters:

Name Type Description Default
predictions torch.Tensor predictions of shape (n_mc x batch_size x n_classes) None

Returns:

Type Description
torch.Tensor entropy of the mean predictive distribution
View Source
    @classmethod
    def predictive_entropy(cls, predictions: torch.Tensor) -> torch.Tensor:
        """
        Calculate the entropy of the mean predictive distribution across the MC samples.

        Args:
            predictions (torch.Tensor): predictions of shape (n_mc x batch_size x n_classes)

        Returns:
            torch.Tensor: entropy of the mean predictive distribution
        """
        return torch.distributions.Categorical(probs=predictions.mean(dim=0)).entropy()

to_json

def to_json(
    inference: torch.Tensor,
    uncertainty: Dict[str, Any] = {},
    **json_kwargs
) -> str

Convert the given inference and uncertainty data to a JSON string.

Parameters:

Name Type Description Default
inference torch.Tensor The inference to convert. None
uncertainty Dict[str, Any] The uncertainty to convert. None
**json_kwargs None Additional keyword arguments to pass to json.dumps. None

Returns:

Type Description
str A JSON string representation of the given inference and uncertainty data.
View Source
    @classmethod
    def to_json(cls, inference: torch.Tensor, uncertainty: Dict[str, Any] = {}, **json_kwargs) -> str:
        """
        Convert the given inference and uncertainty data to a JSON string.

        Args:
            inference (torch.Tensor): The inference to convert.
            uncertainty (Dict[str, Any]): The uncertainty to convert.
            **json_kwargs: Additional keyword arguments to pass to `json.dumps`.

        Returns:
            str: A JSON string representation of the given inference and uncertainty data.
        """
        def prepare(v):
            if isinstance(v, torch.Tensor):
                return v.cpu().tolist()
            if isinstance(v, np.ndarray):  # cspell:ignore ndarray
                return v.tolist()
            if isinstance(v, dict):
                return {k: prepare(_v) for k, _v in v.items()}
            return v

        return json.dumps({
            "inference": inference.tolist(),
            "uncertainty": prepare(uncertainty) if uncertainty else {},
        }, **json_kwargs)

MCDropout

class MCDropout(
    /,
    *args,
    **kwargs
)

Monte Carlo (MC) Dropout Uncertainty Estimation

Requirements:

  • model with dropout layers
  • T, number of samples per input (number of monte-carlo samples/forward passes)

References:

View Source
class MCDropout(UncertaintyBase):
    """
    Monte Carlo (MC) Dropout Uncertainty Estimation

    Requirements:

    - model with dropout layers
    - T, number of samples per input (number of monte-carlo samples/forward passes)

    References:

    - Paper: Understanding Measures of Uncertainty for Adversarial Example Detection
             <https://arxiv.org/abs/1803.08533>
    - Code inspiration: <https://github.com/lsgos/uncertainty-adversarial-paper/tree/master>
    """

    @classmethod
    def prediction(cls, input: Tensor, model: Model) -> Tuple[torch.Tensor, Dict[str, Any]]:
        options = cls.get_options(model)
        N = options.get("N", 10)
        softmax = options.get("softmax", False)

        net: Module = model.get_torch_model()
        net.eval()
        set_dropout(net, state=True)

        out_list = []
        for _ in range(N):
            output = net(input).detach()
            # convert to probabilities if necessary
            if softmax:
                output = torch.softmax(output, dim=1)
            out_list.append(output)
        out = torch.stack(out_list, dim=0)  # (n_mc, batch_size, n_classes)

        inference = out.mean(dim=0)
        uncertainty = cls.interpret(out)
        return inference, uncertainty

Ancestors (in MRO)

  • fl_server_ai.uncertainty.UncertaintyBase
  • abc.ABC

Static methods

expected_entropy

def expected_entropy(
    predictions: torch.Tensor
) -> torch.Tensor

Calculate the mean entropy of the predictive distribution across the MC samples.

Parameters:

Name Type Description Default
predictions torch.Tensor predictions of shape (n_mc x batch_size x n_classes) None

Returns:

Type Description
torch.Tensor mean entropy of the predictive distribution
View Source
    @classmethod
    def expected_entropy(cls, predictions: torch.Tensor) -> torch.Tensor:
        """
        Calculate the mean entropy of the predictive distribution across the MC samples.

        Args:
            predictions (torch.Tensor): predictions of shape (n_mc x batch_size x n_classes)

        Returns:
            torch.Tensor: mean entropy of the predictive distribution
        """
        return torch.distributions.Categorical(probs=predictions).entropy().mean(dim=0)

get_options

def get_options(
    obj: fl_server_core.models.model.Model | fl_server_core.models.training.Training
) -> Dict[str, Any]

Get uncertainty options from training options.

Parameters:

Name Type Description Default
obj Model Training The Model or Training object to retrieve options for.

Returns:

Type Description
Dict[str, Any] Uncertainty options.

Raises:

Type Description
TypeError If the given object is not a Model or Training.
View Source
    @classmethod
    def get_options(cls, obj: Model | Training) -> Dict[str, Any]:
        """
        Get uncertainty options from training options.

        Args:
            obj (Model | Training): The Model or Training object to retrieve options for.

        Returns:
            Dict[str, Any]: Uncertainty options.

        Raises:
            TypeError: If the given object is not a Model or Training.
        """
        if isinstance(obj, Model):
            return Training.objects.filter(model=obj) \
                .values("options") \
                .first()["options"] \
                .get("uncertainty", {})
        if isinstance(obj, Training):
            return obj.options.get("uncertainty", {})
        raise TypeError(f"Expected Model or Training, got {type(obj)}")

interpret

def interpret(
    outputs: torch.Tensor
) -> Dict[str, Any]

Interpret the different network (model) outputs and calculate the uncertainty.

Parameters:

Name Type Description Default
outputs torch.Tensor multiple network (model) outputs (N, batch_size, n_classes) None
View Source
    @classmethod
    def interpret(cls, outputs: torch.Tensor) -> Dict[str, Any]:
        """
        Interpret the different network (model) outputs and calculate the uncertainty.

        Args:
            outputs (torch.Tensor): multiple network (model) outputs (N, batch_size, n_classes)

        Return:
            Tuple[torch.Tensor, Dict[str, Any]]: inference and uncertainty
        """
        variance = outputs.var(dim=0)
        std = outputs.std(dim=0)
        if not (torch.all(outputs <= 1.) and torch.all(outputs >= 0.)):
            return dict(variance=variance, std=std)

        predictive_entropy = cls.predictive_entropy(outputs)
        expected_entropy = cls.expected_entropy(outputs)
        mutual_info = predictive_entropy - expected_entropy  # see cls.mutual_information
        return dict(
            variance=variance,
            std=std,
            predictive_entropy=predictive_entropy,
            expected_entropy=expected_entropy,
            mutual_info=mutual_info,
        )

mutual_information

def mutual_information(
    predictions: torch.Tensor
) -> torch.Tensor

Calculate the BALD (Bayesian Active Learning by Disagreement) of a model;

the difference between the mean of the entropy and the entropy of the mean of the predicted distribution on the predictions. This method is also sometimes referred to as the mutual information (MI).

Parameters:

Name Type Description Default
predictions torch.Tensor predictions of shape (n_mc x batch_size x n_classes) None

Returns:

Type Description
torch.Tensor difference between the mean of the entropy and the entropy of the mean
of the predicted distribution
View Source
    @classmethod
    def mutual_information(cls, predictions: torch.Tensor) -> torch.Tensor:
        """
        Calculate the BALD (Bayesian Active Learning by Disagreement) of a model;
        the difference between the mean of the entropy and the entropy of the mean
        of the predicted distribution on the predictions.
        This method is also sometimes referred to as the mutual information (MI).

        Args:
            predictions (torch.Tensor): predictions of shape (n_mc x batch_size x n_classes)

        Returns:
            torch.Tensor: difference between the mean of the entropy and the entropy of the mean
                    of the predicted distribution
        """
        return cls.predictive_entropy(predictions) - cls.expected_entropy(predictions)

prediction

def prediction(
    input: torch.Tensor,
    model: fl_server_core.models.model.Model
) -> Tuple[torch.Tensor, Dict[str, Any]]

Make a prediction using the given input and model.

Parameters:

Name Type Description Default
input torch.Tensor The input to make a prediction for. None
model Model The model to use for making the prediction. None

Returns:

Type Description
Tuple[torch.Tensor, Dict[str, Any]] The prediction and any associated uncertainty.
View Source
    @classmethod
    def prediction(cls, input: Tensor, model: Model) -> Tuple[torch.Tensor, Dict[str, Any]]:
        options = cls.get_options(model)
        N = options.get("N", 10)
        softmax = options.get("softmax", False)

        net: Module = model.get_torch_model()
        net.eval()
        set_dropout(net, state=True)

        out_list = []
        for _ in range(N):
            output = net(input).detach()
            # convert to probabilities if necessary
            if softmax:
                output = torch.softmax(output, dim=1)
            out_list.append(output)
        out = torch.stack(out_list, dim=0)  # (n_mc, batch_size, n_classes)

        inference = out.mean(dim=0)
        uncertainty = cls.interpret(out)
        return inference, uncertainty

predictive_entropy

def predictive_entropy(
    predictions: torch.Tensor
) -> torch.Tensor

Calculate the entropy of the mean predictive distribution across the MC samples.

Parameters:

Name Type Description Default
predictions torch.Tensor predictions of shape (n_mc x batch_size x n_classes) None

Returns:

Type Description
torch.Tensor entropy of the mean predictive distribution
View Source
    @classmethod
    def predictive_entropy(cls, predictions: torch.Tensor) -> torch.Tensor:
        """
        Calculate the entropy of the mean predictive distribution across the MC samples.

        Args:
            predictions (torch.Tensor): predictions of shape (n_mc x batch_size x n_classes)

        Returns:
            torch.Tensor: entropy of the mean predictive distribution
        """
        return torch.distributions.Categorical(probs=predictions.mean(dim=0)).entropy()

to_json

def to_json(
    inference: torch.Tensor,
    uncertainty: Dict[str, Any] = {},
    **json_kwargs
) -> str

Convert the given inference and uncertainty data to a JSON string.

Parameters:

Name Type Description Default
inference torch.Tensor The inference to convert. None
uncertainty Dict[str, Any] The uncertainty to convert. None
**json_kwargs None Additional keyword arguments to pass to json.dumps. None

Returns:

Type Description
str A JSON string representation of the given inference and uncertainty data.
View Source
    @classmethod
    def to_json(cls, inference: torch.Tensor, uncertainty: Dict[str, Any] = {}, **json_kwargs) -> str:
        """
        Convert the given inference and uncertainty data to a JSON string.

        Args:
            inference (torch.Tensor): The inference to convert.
            uncertainty (Dict[str, Any]): The uncertainty to convert.
            **json_kwargs: Additional keyword arguments to pass to `json.dumps`.

        Returns:
            str: A JSON string representation of the given inference and uncertainty data.
        """
        def prepare(v):
            if isinstance(v, torch.Tensor):
                return v.cpu().tolist()
            if isinstance(v, np.ndarray):  # cspell:ignore ndarray
                return v.tolist()
            if isinstance(v, dict):
                return {k: prepare(_v) for k, _v in v.items()}
            return v

        return json.dumps({
            "inference": inference.tolist(),
            "uncertainty": prepare(uncertainty) if uncertainty else {},
        }, **json_kwargs)

NoneUncertainty

class NoneUncertainty(
    /,
    *args,
    **kwargs
)

Empty uncertainty estimation when no specific uncertainty method is used.

This class does not calculate any uncertainty and only returns the prediction with an empty uncertainty dictionary.

View Source
class NoneUncertainty(UncertaintyBase):
    """
    Empty uncertainty estimation when no specific uncertainty method is used.

    This class does not calculate any uncertainty and only returns the prediction with an empty uncertainty dictionary.
    """

    @classmethod
    def prediction(cls, input: torch.Tensor, model: Model) -> Tuple[torch.Tensor, Dict[str, Any]]:
        net: torch.nn.Module = model.get_torch_model()
        prediction: torch.Tensor = net(input)
        return prediction.argmax(dim=1), {}

Ancestors (in MRO)

  • fl_server_ai.uncertainty.UncertaintyBase
  • abc.ABC

Static methods

expected_entropy

def expected_entropy(
    predictions: torch.Tensor
) -> torch.Tensor

Calculate the mean entropy of the predictive distribution across the MC samples.

Parameters:

Name Type Description Default
predictions torch.Tensor predictions of shape (n_mc x batch_size x n_classes) None

Returns:

Type Description
torch.Tensor mean entropy of the predictive distribution
View Source
    @classmethod
    def expected_entropy(cls, predictions: torch.Tensor) -> torch.Tensor:
        """
        Calculate the mean entropy of the predictive distribution across the MC samples.

        Args:
            predictions (torch.Tensor): predictions of shape (n_mc x batch_size x n_classes)

        Returns:
            torch.Tensor: mean entropy of the predictive distribution
        """
        return torch.distributions.Categorical(probs=predictions).entropy().mean(dim=0)

get_options

def get_options(
    obj: fl_server_core.models.model.Model | fl_server_core.models.training.Training
) -> Dict[str, Any]

Get uncertainty options from training options.

Parameters:

Name Type Description Default
obj Model Training The Model or Training object to retrieve options for.

Returns:

Type Description
Dict[str, Any] Uncertainty options.

Raises:

Type Description
TypeError If the given object is not a Model or Training.
View Source
    @classmethod
    def get_options(cls, obj: Model | Training) -> Dict[str, Any]:
        """
        Get uncertainty options from training options.

        Args:
            obj (Model | Training): The Model or Training object to retrieve options for.

        Returns:
            Dict[str, Any]: Uncertainty options.

        Raises:
            TypeError: If the given object is not a Model or Training.
        """
        if isinstance(obj, Model):
            return Training.objects.filter(model=obj) \
                .values("options") \
                .first()["options"] \
                .get("uncertainty", {})
        if isinstance(obj, Training):
            return obj.options.get("uncertainty", {})
        raise TypeError(f"Expected Model or Training, got {type(obj)}")

interpret

def interpret(
    outputs: torch.Tensor
) -> Dict[str, Any]

Interpret the different network (model) outputs and calculate the uncertainty.

Parameters:

Name Type Description Default
outputs torch.Tensor multiple network (model) outputs (N, batch_size, n_classes) None
View Source
    @classmethod
    def interpret(cls, outputs: torch.Tensor) -> Dict[str, Any]:
        """
        Interpret the different network (model) outputs and calculate the uncertainty.

        Args:
            outputs (torch.Tensor): multiple network (model) outputs (N, batch_size, n_classes)

        Return:
            Tuple[torch.Tensor, Dict[str, Any]]: inference and uncertainty
        """
        variance = outputs.var(dim=0)
        std = outputs.std(dim=0)
        if not (torch.all(outputs <= 1.) and torch.all(outputs >= 0.)):
            return dict(variance=variance, std=std)

        predictive_entropy = cls.predictive_entropy(outputs)
        expected_entropy = cls.expected_entropy(outputs)
        mutual_info = predictive_entropy - expected_entropy  # see cls.mutual_information
        return dict(
            variance=variance,
            std=std,
            predictive_entropy=predictive_entropy,
            expected_entropy=expected_entropy,
            mutual_info=mutual_info,
        )

mutual_information

def mutual_information(
    predictions: torch.Tensor
) -> torch.Tensor

Calculate the BALD (Bayesian Active Learning by Disagreement) of a model;

the difference between the mean of the entropy and the entropy of the mean of the predicted distribution on the predictions. This method is also sometimes referred to as the mutual information (MI).

Parameters:

Name Type Description Default
predictions torch.Tensor predictions of shape (n_mc x batch_size x n_classes) None

Returns:

Type Description
torch.Tensor difference between the mean of the entropy and the entropy of the mean
of the predicted distribution
View Source
    @classmethod
    def mutual_information(cls, predictions: torch.Tensor) -> torch.Tensor:
        """
        Calculate the BALD (Bayesian Active Learning by Disagreement) of a model;
        the difference between the mean of the entropy and the entropy of the mean
        of the predicted distribution on the predictions.
        This method is also sometimes referred to as the mutual information (MI).

        Args:
            predictions (torch.Tensor): predictions of shape (n_mc x batch_size x n_classes)

        Returns:
            torch.Tensor: difference between the mean of the entropy and the entropy of the mean
                    of the predicted distribution
        """
        return cls.predictive_entropy(predictions) - cls.expected_entropy(predictions)

prediction

def prediction(
    input: torch.Tensor,
    model: fl_server_core.models.model.Model
) -> Tuple[torch.Tensor, Dict[str, Any]]

Make a prediction using the given input and model.

Parameters:

Name Type Description Default
input torch.Tensor The input to make a prediction for. None
model Model The model to use for making the prediction. None

Returns:

Type Description
Tuple[torch.Tensor, Dict[str, Any]] The prediction and any associated uncertainty.
View Source
    @classmethod
    def prediction(cls, input: torch.Tensor, model: Model) -> Tuple[torch.Tensor, Dict[str, Any]]:
        net: torch.nn.Module = model.get_torch_model()
        prediction: torch.Tensor = net(input)
        return prediction.argmax(dim=1), {}

predictive_entropy

def predictive_entropy(
    predictions: torch.Tensor
) -> torch.Tensor

Calculate the entropy of the mean predictive distribution across the MC samples.

Parameters:

Name Type Description Default
predictions torch.Tensor predictions of shape (n_mc x batch_size x n_classes) None

Returns:

Type Description
torch.Tensor entropy of the mean predictive distribution
View Source
    @classmethod
    def predictive_entropy(cls, predictions: torch.Tensor) -> torch.Tensor:
        """
        Calculate the entropy of the mean predictive distribution across the MC samples.

        Args:
            predictions (torch.Tensor): predictions of shape (n_mc x batch_size x n_classes)

        Returns:
            torch.Tensor: entropy of the mean predictive distribution
        """
        return torch.distributions.Categorical(probs=predictions.mean(dim=0)).entropy()

to_json

def to_json(
    inference: torch.Tensor,
    uncertainty: Dict[str, Any] = {},
    **json_kwargs
) -> str

Convert the given inference and uncertainty data to a JSON string.

Parameters:

Name Type Description Default
inference torch.Tensor The inference to convert. None
uncertainty Dict[str, Any] The uncertainty to convert. None
**json_kwargs None Additional keyword arguments to pass to json.dumps. None

Returns:

Type Description
str A JSON string representation of the given inference and uncertainty data.
View Source
    @classmethod
    def to_json(cls, inference: torch.Tensor, uncertainty: Dict[str, Any] = {}, **json_kwargs) -> str:
        """
        Convert the given inference and uncertainty data to a JSON string.

        Args:
            inference (torch.Tensor): The inference to convert.
            uncertainty (Dict[str, Any]): The uncertainty to convert.
            **json_kwargs: Additional keyword arguments to pass to `json.dumps`.

        Returns:
            str: A JSON string representation of the given inference and uncertainty data.
        """
        def prepare(v):
            if isinstance(v, torch.Tensor):
                return v.cpu().tolist()
            if isinstance(v, np.ndarray):  # cspell:ignore ndarray
                return v.tolist()
            if isinstance(v, dict):
                return {k: prepare(_v) for k, _v in v.items()}
            return v

        return json.dumps({
            "inference": inference.tolist(),
            "uncertainty": prepare(uncertainty) if uncertainty else {},
        }, **json_kwargs)

SWAG

class SWAG(
    /,
    *args,
    **kwargs
)

Stochastic Weight Averaging Gaussian (SWAG) uncertainty estimation.

View Source
class SWAG(UncertaintyBase):
    """
    Stochastic Weight Averaging Gaussian (SWAG) uncertainty estimation.
    """

    @classmethod
    def prediction(cls, input: torch.Tensor, model: SWAGModel) -> Tuple[torch.Tensor, Dict[str, Any]]:
        options = cls.get_options(model)
        N = options.get("N", 10)

        net: torch.nn.Module = model.get_torch_model()

        # first and second moment are already ensured to be in
        # alphabetical order in the database
        fm = model.first_moment
        sm = model.second_moment
        std = sm - torch.pow(fm, 2)
        params = torch.normal(mean=fm[None, :], std=std).expand(N, -1)

        prediction_list = []
        for n in range(N):
            torch.nn.utils.vector_to_parameters(params[n], net.parameters())
            prediction = net(input)
            prediction_list.append(prediction)
        predictions = torch.stack(prediction_list)

        inference = predictions.mean(dim=0)
        uncertainty = cls.interpret(predictions)
        return inference, uncertainty

Ancestors (in MRO)

  • fl_server_ai.uncertainty.UncertaintyBase
  • abc.ABC

Static methods

expected_entropy

def expected_entropy(
    predictions: torch.Tensor
) -> torch.Tensor

Calculate the mean entropy of the predictive distribution across the MC samples.

Parameters:

Name Type Description Default
predictions torch.Tensor predictions of shape (n_mc x batch_size x n_classes) None

Returns:

Type Description
torch.Tensor mean entropy of the predictive distribution
View Source
    @classmethod
    def expected_entropy(cls, predictions: torch.Tensor) -> torch.Tensor:
        """
        Calculate the mean entropy of the predictive distribution across the MC samples.

        Args:
            predictions (torch.Tensor): predictions of shape (n_mc x batch_size x n_classes)

        Returns:
            torch.Tensor: mean entropy of the predictive distribution
        """
        return torch.distributions.Categorical(probs=predictions).entropy().mean(dim=0)

get_options

def get_options(
    obj: fl_server_core.models.model.Model | fl_server_core.models.training.Training
) -> Dict[str, Any]

Get uncertainty options from training options.

Parameters:

Name Type Description Default
obj Model Training The Model or Training object to retrieve options for.

Returns:

Type Description
Dict[str, Any] Uncertainty options.

Raises:

Type Description
TypeError If the given object is not a Model or Training.
View Source
    @classmethod
    def get_options(cls, obj: Model | Training) -> Dict[str, Any]:
        """
        Get uncertainty options from training options.

        Args:
            obj (Model | Training): The Model or Training object to retrieve options for.

        Returns:
            Dict[str, Any]: Uncertainty options.

        Raises:
            TypeError: If the given object is not a Model or Training.
        """
        if isinstance(obj, Model):
            return Training.objects.filter(model=obj) \
                .values("options") \
                .first()["options"] \
                .get("uncertainty", {})
        if isinstance(obj, Training):
            return obj.options.get("uncertainty", {})
        raise TypeError(f"Expected Model or Training, got {type(obj)}")

interpret

def interpret(
    outputs: torch.Tensor
) -> Dict[str, Any]

Interpret the different network (model) outputs and calculate the uncertainty.

Parameters:

Name Type Description Default
outputs torch.Tensor multiple network (model) outputs (N, batch_size, n_classes) None
View Source
    @classmethod
    def interpret(cls, outputs: torch.Tensor) -> Dict[str, Any]:
        """
        Interpret the different network (model) outputs and calculate the uncertainty.

        Args:
            outputs (torch.Tensor): multiple network (model) outputs (N, batch_size, n_classes)

        Return:
            Tuple[torch.Tensor, Dict[str, Any]]: inference and uncertainty
        """
        variance = outputs.var(dim=0)
        std = outputs.std(dim=0)
        if not (torch.all(outputs <= 1.) and torch.all(outputs >= 0.)):
            return dict(variance=variance, std=std)

        predictive_entropy = cls.predictive_entropy(outputs)
        expected_entropy = cls.expected_entropy(outputs)
        mutual_info = predictive_entropy - expected_entropy  # see cls.mutual_information
        return dict(
            variance=variance,
            std=std,
            predictive_entropy=predictive_entropy,
            expected_entropy=expected_entropy,
            mutual_info=mutual_info,
        )

mutual_information

def mutual_information(
    predictions: torch.Tensor
) -> torch.Tensor

Calculate the BALD (Bayesian Active Learning by Disagreement) of a model;

the difference between the mean of the entropy and the entropy of the mean of the predicted distribution on the predictions. This method is also sometimes referred to as the mutual information (MI).

Parameters:

Name Type Description Default
predictions torch.Tensor predictions of shape (n_mc x batch_size x n_classes) None

Returns:

Type Description
torch.Tensor difference between the mean of the entropy and the entropy of the mean
of the predicted distribution
View Source
    @classmethod
    def mutual_information(cls, predictions: torch.Tensor) -> torch.Tensor:
        """
        Calculate the BALD (Bayesian Active Learning by Disagreement) of a model;
        the difference between the mean of the entropy and the entropy of the mean
        of the predicted distribution on the predictions.
        This method is also sometimes referred to as the mutual information (MI).

        Args:
            predictions (torch.Tensor): predictions of shape (n_mc x batch_size x n_classes)

        Returns:
            torch.Tensor: difference between the mean of the entropy and the entropy of the mean
                    of the predicted distribution
        """
        return cls.predictive_entropy(predictions) - cls.expected_entropy(predictions)

prediction

def prediction(
    input: torch.Tensor,
    model: fl_server_core.models.model.SWAGModel
) -> Tuple[torch.Tensor, Dict[str, Any]]

Make a prediction using the given input and model.

Parameters:

Name Type Description Default
input torch.Tensor The input to make a prediction for. None
model Model The model to use for making the prediction. None

Returns:

Type Description
Tuple[torch.Tensor, Dict[str, Any]] The prediction and any associated uncertainty.
View Source
    @classmethod
    def prediction(cls, input: torch.Tensor, model: SWAGModel) -> Tuple[torch.Tensor, Dict[str, Any]]:
        options = cls.get_options(model)
        N = options.get("N", 10)

        net: torch.nn.Module = model.get_torch_model()

        # first and second moment are already ensured to be in
        # alphabetical order in the database
        fm = model.first_moment
        sm = model.second_moment
        std = sm - torch.pow(fm, 2)
        params = torch.normal(mean=fm[None, :], std=std).expand(N, -1)

        prediction_list = []
        for n in range(N):
            torch.nn.utils.vector_to_parameters(params[n], net.parameters())
            prediction = net(input)
            prediction_list.append(prediction)
        predictions = torch.stack(prediction_list)

        inference = predictions.mean(dim=0)
        uncertainty = cls.interpret(predictions)
        return inference, uncertainty

predictive_entropy

def predictive_entropy(
    predictions: torch.Tensor
) -> torch.Tensor

Calculate the entropy of the mean predictive distribution across the MC samples.

Parameters:

Name Type Description Default
predictions torch.Tensor predictions of shape (n_mc x batch_size x n_classes) None

Returns:

Type Description
torch.Tensor entropy of the mean predictive distribution
View Source
    @classmethod
    def predictive_entropy(cls, predictions: torch.Tensor) -> torch.Tensor:
        """
        Calculate the entropy of the mean predictive distribution across the MC samples.

        Args:
            predictions (torch.Tensor): predictions of shape (n_mc x batch_size x n_classes)

        Returns:
            torch.Tensor: entropy of the mean predictive distribution
        """
        return torch.distributions.Categorical(probs=predictions.mean(dim=0)).entropy()

to_json

def to_json(
    inference: torch.Tensor,
    uncertainty: Dict[str, Any] = {},
    **json_kwargs
) -> str

Convert the given inference and uncertainty data to a JSON string.

Parameters:

Name Type Description Default
inference torch.Tensor The inference to convert. None
uncertainty Dict[str, Any] The uncertainty to convert. None
**json_kwargs None Additional keyword arguments to pass to json.dumps. None

Returns:

Type Description
str A JSON string representation of the given inference and uncertainty data.
View Source
    @classmethod
    def to_json(cls, inference: torch.Tensor, uncertainty: Dict[str, Any] = {}, **json_kwargs) -> str:
        """
        Convert the given inference and uncertainty data to a JSON string.

        Args:
            inference (torch.Tensor): The inference to convert.
            uncertainty (Dict[str, Any]): The uncertainty to convert.
            **json_kwargs: Additional keyword arguments to pass to `json.dumps`.

        Returns:
            str: A JSON string representation of the given inference and uncertainty data.
        """
        def prepare(v):
            if isinstance(v, torch.Tensor):
                return v.cpu().tolist()
            if isinstance(v, np.ndarray):  # cspell:ignore ndarray
                return v.tolist()
            if isinstance(v, dict):
                return {k: prepare(_v) for k, _v in v.items()}
            return v

        return json.dumps({
            "inference": inference.tolist(),
            "uncertainty": prepare(uncertainty) if uncertainty else {},
        }, **json_kwargs)

UncertaintyBase

class UncertaintyBase(
    /,
    *args,
    **kwargs
)

Abstract base class for uncertainty estimation.

This class defines the interface for uncertainty estimation in federated learning.

View Source
class UncertaintyBase(ABC):
    """
    Abstract base class for uncertainty estimation.

    This class defines the interface for uncertainty estimation in federated learning.
    """

    _logger = getLogger("fl.server")
    """The private logger instance for the uncertainty estimation."""

    @classmethod
    @abstractmethod
    def prediction(cls, input: torch.Tensor, model: Model) -> Tuple[torch.Tensor, Dict[str, Any]]:
        """
        Make a prediction using the given input and model.

        Args:
            input (torch.Tensor): The input to make a prediction for.
            model (Model): The model to use for making the prediction.

        Returns:
            Tuple[torch.Tensor, Dict[str, Any]]: The prediction and any associated uncertainty.
        """
        pass

    @classmethod
    def interpret(cls, outputs: torch.Tensor) -> Dict[str, Any]:
        """
        Interpret the different network (model) outputs and calculate the uncertainty.

        Args:
            outputs (torch.Tensor): multiple network (model) outputs (N, batch_size, n_classes)

        Return:
            Tuple[torch.Tensor, Dict[str, Any]]: inference and uncertainty
        """
        variance = outputs.var(dim=0)
        std = outputs.std(dim=0)
        if not (torch.all(outputs <= 1.) and torch.all(outputs >= 0.)):
            return dict(variance=variance, std=std)

        predictive_entropy = cls.predictive_entropy(outputs)
        expected_entropy = cls.expected_entropy(outputs)
        mutual_info = predictive_entropy - expected_entropy  # see cls.mutual_information
        return dict(
            variance=variance,
            std=std,
            predictive_entropy=predictive_entropy,
            expected_entropy=expected_entropy,
            mutual_info=mutual_info,
        )

    @classmethod
    def expected_entropy(cls, predictions: torch.Tensor) -> torch.Tensor:
        """
        Calculate the mean entropy of the predictive distribution across the MC samples.

        Args:
            predictions (torch.Tensor): predictions of shape (n_mc x batch_size x n_classes)

        Returns:
            torch.Tensor: mean entropy of the predictive distribution
        """
        return torch.distributions.Categorical(probs=predictions).entropy().mean(dim=0)

    @classmethod
    def predictive_entropy(cls, predictions: torch.Tensor) -> torch.Tensor:
        """
        Calculate the entropy of the mean predictive distribution across the MC samples.

        Args:
            predictions (torch.Tensor): predictions of shape (n_mc x batch_size x n_classes)

        Returns:
            torch.Tensor: entropy of the mean predictive distribution
        """
        return torch.distributions.Categorical(probs=predictions.mean(dim=0)).entropy()

    @classmethod
    def mutual_information(cls, predictions: torch.Tensor) -> torch.Tensor:
        """
        Calculate the BALD (Bayesian Active Learning by Disagreement) of a model;
        the difference between the mean of the entropy and the entropy of the mean
        of the predicted distribution on the predictions.
        This method is also sometimes referred to as the mutual information (MI).

        Args:
            predictions (torch.Tensor): predictions of shape (n_mc x batch_size x n_classes)

        Returns:
            torch.Tensor: difference between the mean of the entropy and the entropy of the mean
                    of the predicted distribution
        """
        return cls.predictive_entropy(predictions) - cls.expected_entropy(predictions)

    @classmethod
    def get_options(cls, obj: Model | Training) -> Dict[str, Any]:
        """
        Get uncertainty options from training options.

        Args:
            obj (Model | Training): The Model or Training object to retrieve options for.

        Returns:
            Dict[str, Any]: Uncertainty options.

        Raises:
            TypeError: If the given object is not a Model or Training.
        """
        if isinstance(obj, Model):
            return Training.objects.filter(model=obj) \
                .values("options") \
                .first()["options"] \
                .get("uncertainty", {})
        if isinstance(obj, Training):
            return obj.options.get("uncertainty", {})
        raise TypeError(f"Expected Model or Training, got {type(obj)}")

    @classmethod
    def to_json(cls, inference: torch.Tensor, uncertainty: Dict[str, Any] = {}, **json_kwargs) -> str:
        """
        Convert the given inference and uncertainty data to a JSON string.

        Args:
            inference (torch.Tensor): The inference to convert.
            uncertainty (Dict[str, Any]): The uncertainty to convert.
            **json_kwargs: Additional keyword arguments to pass to `json.dumps`.

        Returns:
            str: A JSON string representation of the given inference and uncertainty data.
        """
        def prepare(v):
            if isinstance(v, torch.Tensor):
                return v.cpu().tolist()
            if isinstance(v, np.ndarray):  # cspell:ignore ndarray
                return v.tolist()
            if isinstance(v, dict):
                return {k: prepare(_v) for k, _v in v.items()}
            return v

        return json.dumps({
            "inference": inference.tolist(),
            "uncertainty": prepare(uncertainty) if uncertainty else {},
        }, **json_kwargs)

Ancestors (in MRO)

  • abc.ABC

Descendants

  • fl_server_ai.uncertainty.Ensemble
  • fl_server_ai.uncertainty.MCDropout
  • fl_server_ai.uncertainty.NoneUncertainty
  • fl_server_ai.uncertainty.SWAG

Static methods

expected_entropy

def expected_entropy(
    predictions: torch.Tensor
) -> torch.Tensor

Calculate the mean entropy of the predictive distribution across the MC samples.

Parameters:

Name Type Description Default
predictions torch.Tensor predictions of shape (n_mc x batch_size x n_classes) None

Returns:

Type Description
torch.Tensor mean entropy of the predictive distribution
View Source
    @classmethod
    def expected_entropy(cls, predictions: torch.Tensor) -> torch.Tensor:
        """
        Calculate the mean entropy of the predictive distribution across the MC samples.

        Args:
            predictions (torch.Tensor): predictions of shape (n_mc x batch_size x n_classes)

        Returns:
            torch.Tensor: mean entropy of the predictive distribution
        """
        return torch.distributions.Categorical(probs=predictions).entropy().mean(dim=0)

get_options

def get_options(
    obj: fl_server_core.models.model.Model | fl_server_core.models.training.Training
) -> Dict[str, Any]

Get uncertainty options from training options.

Parameters:

Name Type Description Default
obj Model Training The Model or Training object to retrieve options for.

Returns:

Type Description
Dict[str, Any] Uncertainty options.

Raises:

Type Description
TypeError If the given object is not a Model or Training.
View Source
    @classmethod
    def get_options(cls, obj: Model | Training) -> Dict[str, Any]:
        """
        Get uncertainty options from training options.

        Args:
            obj (Model | Training): The Model or Training object to retrieve options for.

        Returns:
            Dict[str, Any]: Uncertainty options.

        Raises:
            TypeError: If the given object is not a Model or Training.
        """
        if isinstance(obj, Model):
            return Training.objects.filter(model=obj) \
                .values("options") \
                .first()["options"] \
                .get("uncertainty", {})
        if isinstance(obj, Training):
            return obj.options.get("uncertainty", {})
        raise TypeError(f"Expected Model or Training, got {type(obj)}")

interpret

def interpret(
    outputs: torch.Tensor
) -> Dict[str, Any]

Interpret the different network (model) outputs and calculate the uncertainty.

Parameters:

Name Type Description Default
outputs torch.Tensor multiple network (model) outputs (N, batch_size, n_classes) None
View Source
    @classmethod
    def interpret(cls, outputs: torch.Tensor) -> Dict[str, Any]:
        """
        Interpret the different network (model) outputs and calculate the uncertainty.

        Args:
            outputs (torch.Tensor): multiple network (model) outputs (N, batch_size, n_classes)

        Return:
            Tuple[torch.Tensor, Dict[str, Any]]: inference and uncertainty
        """
        variance = outputs.var(dim=0)
        std = outputs.std(dim=0)
        if not (torch.all(outputs <= 1.) and torch.all(outputs >= 0.)):
            return dict(variance=variance, std=std)

        predictive_entropy = cls.predictive_entropy(outputs)
        expected_entropy = cls.expected_entropy(outputs)
        mutual_info = predictive_entropy - expected_entropy  # see cls.mutual_information
        return dict(
            variance=variance,
            std=std,
            predictive_entropy=predictive_entropy,
            expected_entropy=expected_entropy,
            mutual_info=mutual_info,
        )

mutual_information

def mutual_information(
    predictions: torch.Tensor
) -> torch.Tensor

Calculate the BALD (Bayesian Active Learning by Disagreement) of a model;

the difference between the mean of the entropy and the entropy of the mean of the predicted distribution on the predictions. This method is also sometimes referred to as the mutual information (MI).

Parameters:

Name Type Description Default
predictions torch.Tensor predictions of shape (n_mc x batch_size x n_classes) None

Returns:

Type Description
torch.Tensor difference between the mean of the entropy and the entropy of the mean
of the predicted distribution
View Source
    @classmethod
    def mutual_information(cls, predictions: torch.Tensor) -> torch.Tensor:
        """
        Calculate the BALD (Bayesian Active Learning by Disagreement) of a model;
        the difference between the mean of the entropy and the entropy of the mean
        of the predicted distribution on the predictions.
        This method is also sometimes referred to as the mutual information (MI).

        Args:
            predictions (torch.Tensor): predictions of shape (n_mc x batch_size x n_classes)

        Returns:
            torch.Tensor: difference between the mean of the entropy and the entropy of the mean
                    of the predicted distribution
        """
        return cls.predictive_entropy(predictions) - cls.expected_entropy(predictions)

prediction

def prediction(
    input: torch.Tensor,
    model: fl_server_core.models.model.Model
) -> Tuple[torch.Tensor, Dict[str, Any]]

Make a prediction using the given input and model.

Parameters:

Name Type Description Default
input torch.Tensor The input to make a prediction for. None
model Model The model to use for making the prediction. None

Returns:

Type Description
Tuple[torch.Tensor, Dict[str, Any]] The prediction and any associated uncertainty.
View Source
    @classmethod
    @abstractmethod
    def prediction(cls, input: torch.Tensor, model: Model) -> Tuple[torch.Tensor, Dict[str, Any]]:
        """
        Make a prediction using the given input and model.

        Args:
            input (torch.Tensor): The input to make a prediction for.
            model (Model): The model to use for making the prediction.

        Returns:
            Tuple[torch.Tensor, Dict[str, Any]]: The prediction and any associated uncertainty.
        """
        pass

predictive_entropy

def predictive_entropy(
    predictions: torch.Tensor
) -> torch.Tensor

Calculate the entropy of the mean predictive distribution across the MC samples.

Parameters:

Name Type Description Default
predictions torch.Tensor predictions of shape (n_mc x batch_size x n_classes) None

Returns:

Type Description
torch.Tensor entropy of the mean predictive distribution
View Source
    @classmethod
    def predictive_entropy(cls, predictions: torch.Tensor) -> torch.Tensor:
        """
        Calculate the entropy of the mean predictive distribution across the MC samples.

        Args:
            predictions (torch.Tensor): predictions of shape (n_mc x batch_size x n_classes)

        Returns:
            torch.Tensor: entropy of the mean predictive distribution
        """
        return torch.distributions.Categorical(probs=predictions.mean(dim=0)).entropy()

to_json

def to_json(
    inference: torch.Tensor,
    uncertainty: Dict[str, Any] = {},
    **json_kwargs
) -> str

Convert the given inference and uncertainty data to a JSON string.

Parameters:

Name Type Description Default
inference torch.Tensor The inference to convert. None
uncertainty Dict[str, Any] The uncertainty to convert. None
**json_kwargs None Additional keyword arguments to pass to json.dumps. None

Returns:

Type Description
str A JSON string representation of the given inference and uncertainty data.
View Source
    @classmethod
    def to_json(cls, inference: torch.Tensor, uncertainty: Dict[str, Any] = {}, **json_kwargs) -> str:
        """
        Convert the given inference and uncertainty data to a JSON string.

        Args:
            inference (torch.Tensor): The inference to convert.
            uncertainty (Dict[str, Any]): The uncertainty to convert.
            **json_kwargs: Additional keyword arguments to pass to `json.dumps`.

        Returns:
            str: A JSON string representation of the given inference and uncertainty data.
        """
        def prepare(v):
            if isinstance(v, torch.Tensor):
                return v.cpu().tolist()
            if isinstance(v, np.ndarray):  # cspell:ignore ndarray
                return v.tolist()
            if isinstance(v, dict):
                return {k: prepare(_v) for k, _v in v.items()}
            return v

        return json.dumps({
            "inference": inference.tolist(),
            "uncertainty": prepare(uncertainty) if uncertainty else {},
        }, **json_kwargs)