core.distribution.AbstractVariable

 1from abc import ABC, abstractmethod
 2
 3import torch
 4import torch.nn as nn
 5
 6from core.utils import KLDivergenceInterface
 7
 8
 9class AbstractVariable(nn.Module, KLDivergenceInterface, ABC):
10    """
11    An abstract class representing a single random variable for a probabilistic
12    neural network parameter (e.g., weight or bias).
13
14    Each variable holds:
15      - A mean parameter (`mu`)
16      - A `rho` parameter that is used to derive the standard deviation (`sigma`)
17      - A method to sample from the underlying distribution
18      - A method to compute KL divergence with another variable of the same type
19
20    This class inherits from `nn.Module` for parameter registration in PyTorch
21    and from `KLDivergenceInterface` for consistent KL divergence handling.
22    """
23
24    def __init__(
25        self,
26        mu: torch.Tensor,
27        rho: torch.Tensor,
28        mu_requires_grad: bool = False,
29        rho_requires_grad: bool = False,
30    ):
31        """
32        Initialize an AbstractVariable with given `mu` and `rho` tensors.
33
34        Args:
35            mu (torch.Tensor): The mean parameter for the distribution.
36            rho (torch.Tensor): The parameter from which we derive sigma = log(1 + exp(rho)).
37            mu_requires_grad (bool, optional): If True, allow gradients on `mu`.
38            rho_requires_grad (bool, optional): If True, allow gradients on `rho`.
39        """
40        super().__init__()
41        self.mu = nn.Parameter(mu.detach().clone(), requires_grad=mu_requires_grad)
42        self.rho = nn.Parameter(rho.detach().clone(), requires_grad=rho_requires_grad)
43        self.kl_div = None
44
45    @property
46    def sigma(self) -> torch.Tensor:
47        """
48        The standard deviation of the distribution, computed as:
49            sigma = log(1 + exp(rho)).
50
51        Returns:
52            torch.Tensor: A tensor representing the current standard deviation.
53        """
54        return torch.log(1 + torch.exp(self.rho))
55
56    @abstractmethod
57    def sample(self) -> torch.Tensor:
58        """
59        Draw a sample from this variable's underlying distribution.
60
61        Returns:
62            torch.Tensor: A sampled value of the same shape as `mu`.
63        """
64        pass
65
66    @abstractmethod
67    def compute_kl(self, other: "AbstractVariable") -> torch.Tensor:
68        """
69        Compute the KL divergence between this variable and another variable
70        of the same distribution type.
71
72        Args:
73            other (AbstractVariable): Another AbstractVariable instance
74                with comparable parameters (e.g., mu, rho).
75
76        Returns:
77            torch.Tensor: A scalar tensor with the KL divergence value.
78        """
79        pass