Source code for fastbnns.bnn.priors

"""Definitions of prior distributions over neural network parameters."""

from collections.abc import Iterable
import copy
from typing import Any, Optional

import torch
import torch.distributions as dist




class Distribution(torch.nn.Module):
    """Distribution wrapper to facilitate device transfers."""

    def __init__(
        self,
        distribution: Optional[dist.Distribution] = None,
        *args: Any,
        **kwargs: Any
    ):
        """Initialize wrapper."""
        super().__init__()

        # If `distribution` is passed, we'll build the wrapper automatically.
        if distribution is not None:
            self._distribution = copy.deepcopy(distribution)
            for key, val in distribution.__dict__.items():
                if isinstance(val, torch.Tensor):
                    self.register_buffer(key, val)

    @property
    def distribution(self) -> dist.Distribution:
        """Prepare an instance of the distribution."""
        for key, val in self._distribution.__dict__.items():
            if isinstance(val, torch.Tensor):
                buffer_val = getattr(self, key)
                setattr(self._distribution, key, buffer_val)
        return self._distribution



    def log_prob(self, x: torch.Tensor) -> torch.Tensor:
        """Compute the log PDF of the prior at points `x`."""
        return self.distribution.log_prob(x)




    def sample(self, sample_shape: Iterable = torch.Size()) -> torch.Tensor:
        """Generate samples from the prior of size `sample_shape`."""
        return self.distribution.sample(sample_shape=sample_shape)






class SpikeSlab(Distribution):
    """Spike-slab Gaussian Mixture Model prior."""

    def __init__(
        self,
        loc: torch.Tensor = torch.tensor([0.0, 0.0]),
        scale: torch.Tensor = torch.tensor([0.1, 1.0]),
        probs: torch.Tensor = torch.tensor([0.5, 0.5]),
    ):
        super().__init__()

        self.register_buffer("loc", loc)
        self.register_buffer("scale", scale)
        self.register_buffer("probs", probs)

    @property
    def distribution(self) -> dist.Distribution:
        """Prepare an instance of the distribution."""
        mixture_distribution = dist.Categorical(probs=self.probs)
        return dist.MixtureSameFamily(
            mixture_distribution=mixture_distribution,
            component_distribution=dist.Normal(loc=self.loc, scale=self.scale),
        )



if __name__ == "__main__":
    import matplotlib.pyplot as plt

    # Example using an SpikeSlab prior.
    prior = SpikeSlab()
    sample = prior.sample(sample_shape=(100, 1))
    x = torch.linspace(-5.0, 5.0, 1000)
    pdf = torch.exp(prior.log_prob(x))
    fig, ax = plt.subplots()
    ax.hist(sample, density=True)
    ax.plot(x, pdf)
    plt.show()