Skip to content

Sampleables

A torch morphism at an input, as the kernel's sampleable.

sites

Host sampleables: a torch morphism applied to an input, as a distribution.

A MonadicProgram step draws from a ContinuousMorphism at an input tensor. The kernel's Random.sample request takes a sampleable value, so the pair is wrapped as one: TorchSampleable offers the sampling and scoring interface the prelude handlers use and keeps the morphism and input a strategy or a trace may want to read.

TorchSampleable dataclass 

TorchSampleable(morphism: ContinuousMorphism, input: Tensor, declared: ContinuousMorphism | None = None)

A morphism applied to an input, offered to the kernel as a sampleable.

PARAMETER DESCRIPTION
morphism

The morphism the site draws from, under whatever parameters the run's Param handlers answered.

TYPE: ContinuousMorphism

input

The input tensor the morphism is conditioned on, carrying the batch axis.

TYPE: Tensor

declared

The morphism as the program declares it, which a trace records and a strategy reads the family of; the drawing morphism itself when no parameter store rebinds it.

TYPE: ContinuousMorphism | None DEFAULT: None

__post_init__ 

__post_init__() -> None

Default the declared morphism to the drawing one.

Source code in src/quivers/effects/sites.py 
41
42
43
44
def __post_init__(self) -> None:
    """Default the declared morphism to the drawing one."""
    if self.declared is None:
        object.__setattr__(self, "declared", self.morphism)

rsample 

rsample() -> Tensor

Draw a reparameterised sample.

RETURNS DESCRIPTION
Tensor

One draw per batch row.
Source code in src/quivers/effects/sites.py 
46
47
48
49
50
51
52
53
54
def rsample(self) -> torch.Tensor:
    """Draw a reparameterised sample.

    Returns
    -------
    torch.Tensor
        One draw per batch row.
    """
    return self.morphism.rsample(self.input)

sample 

sample() -> Tensor

Draw a sample.

RETURNS DESCRIPTION
Tensor

One draw per batch row.
Source code in src/quivers/effects/sites.py 
56
57
58
59
60
61
62
63
64
def sample(self) -> torch.Tensor:
    """Draw a sample.

    Returns
    -------
    torch.Tensor
        One draw per batch row.
    """
    return self.morphism.rsample(self.input)

log_prob 

log_prob(value: Tensor) -> Tensor

Score a value under the morphism at the input.

PARAMETER DESCRIPTION
value

The value, one per batch row.

TYPE: Tensor

RETURNS DESCRIPTION
Tensor

The log density per batch row.
Source code in src/quivers/effects/sites.py 
66
67
68
69
70
71
72
73
74
75
76
77
78
79
def log_prob(self, value: torch.Tensor) -> torch.Tensor:
    """Score a value under the morphism at the input.

    Parameters
    ----------
    value : torch.Tensor
        The value, one per batch row.

    Returns
    -------
    torch.Tensor
        The log density per batch row.
    """
    return self.morphism.log_prob(self.input, value)

parameters 

parameters() -> dict[str, Tensor]

The distribution parameters the morphism computes at the input.

RETURNS DESCRIPTION
dict[str, Tensor]

The named parameters of the torch distribution the morphism builds, or loc and scale for a morphism that exposes only a location-scale pair.
RAISES DESCRIPTION
TypeError

If the morphism exposes neither a distribution builder nor a location-scale pair.
Source code in src/quivers/effects/sites.py 
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
def parameters(self) -> dict[str, torch.Tensor]:
    """The distribution parameters the morphism computes at the input.

    Returns
    -------
    dict[str, torch.Tensor]
        The named parameters of the torch distribution the morphism
        builds, or ``loc`` and ``scale`` for a morphism that exposes
        only a location-scale pair.

    Raises
    ------
    TypeError
        If the morphism exposes neither a distribution builder nor a
        location-scale pair.
    """
    get_dist = getattr(self.morphism, "_get_dist", None)
    if callable(get_dist):
        distribution = get_dist(self.input)
        return {
            name: getattr(distribution, name)
            for name in distribution.arg_constraints
            if hasattr(distribution, name)
        }
    get_params = getattr(self.morphism, "_get_params", None)
    if callable(get_params):
        loc, scale = get_params(self.input)
        return {"loc": loc, "scale": scale}
    raise TypeError(
        f"{type(self.morphism).__name__} exposes no distribution parameters"
    )