torch_mist.models

Submodules

Package Contents

Classes

VIB

MIB

CEB

class torch_mist.models.VIB(mi_estimator: torch_mist.estimators.MIEstimator, q_ZX: torch.distributions.Distribution, beta: float, p_ZX_given_X: torch.nn.Module | pyro.distributions.ConditionalDistribution | None = None, p_ZY_given_Y: torch.nn.Module | pyro.distributions.ConditionalDistribution | None = None)

Bases: InformationBottleneck

regularization(zx: torch.Tensor, zy: torch.Tensor, p_ZX_given_x: torch.distributions.Distribution, p_ZY_given_y: torch.distributions.Distribution)
class torch_mist.models.MIB(mi_estimator: torch_mist.estimators.MIEstimator, beta: float, p_ZX_given_X: torch.nn.Module | pyro.distributions.ConditionalDistribution | None = None, p_ZY_given_Y: torch.nn.Module | pyro.distributions.ConditionalDistribution | None = None)

Bases: InformationBottleneck

regularization(zx: torch.Tensor, zy: torch.Tensor, p_ZX_given_x: torch.distributions.Distribution, p_ZY_given_y: torch.distributions.Distribution)
class torch_mist.models.CEB(mi_estimator: torch_mist.estimators.MIEstimator, q_ZX_given_ZY: pyro.distributions.ConditionalDistribution, beta: float, p_ZX_given_X: torch.nn.Module | pyro.distributions.ConditionalDistribution | None = None, p_ZY_given_Y: torch.nn.Module | pyro.distributions.ConditionalDistribution | None = None)

Bases: InformationBottleneck

regularization(zx: torch.Tensor, zy: torch.Tensor, p_ZX_given_x: torch.distributions.Distribution, p_ZY_given_y: torch.distributions.Distribution)