Boundaries

Boundary conditions and constraints for continuous spaces.

boundaries

Boundary morphisms between discrete and continuous spaces.

These morphisms bridge the gap between the finite tensor world (FinSet, discrete Morphism) and the continuous sampling world (ContinuousSpace, ContinuousMorphism).

Morphisms provided

Discretize — continuous space -> finite set (binning) Embed — finite set -> continuous space (kernel density)

Discretize 

Discretize(domain: Euclidean, n_bins: int, soft: bool = True, temperature: float = 0.1)

Bases: ContinuousMorphism

Map a continuous space to a finite set by binning.

Divides a bounded continuous space into n_bins equal-width bins and assigns each continuous input to the bin containing it. The resulting distribution is deterministic (one-hot on the correct bin).

This is useful for converting continuous outputs into discrete categories for downstream processing in the finite tensor world.

For log_prob: returns log(1) = 0 if y equals the correct bin, log(0) otherwise. In practice uses a soft assignment based on distance to bin centers for gradient flow.

PARAMETER DESCRIPTION
domain

Source continuous space (must be bounded, 1-dimensional).

TYPE: Euclidean

n_bins

Number of discrete bins.

TYPE: int

soft

If True (default), use soft binning via softmax over negative squared distances. If False, use hard (argmax) assignment.

TYPE: bool DEFAULT: True

temperature

Temperature for soft binning. Lower = sharper.

TYPE: float DEFAULT: 0.1

Source code in src/quivers/continuous/boundaries.py 
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
def __init__(
    self,
    domain: Euclidean,
    n_bins: int,
    soft: bool = True,
    temperature: float = 0.1,
) -> None:
    if not isinstance(domain, Euclidean) or not domain.is_bounded:
        raise ValueError("Discretize requires a bounded Euclidean domain")

    if domain.dim != 1:
        raise ValueError(
            f"Discretize currently supports 1-d spaces only, got dim={domain.dim}"
        )

    codomain = FinSet(name="bins", cardinality=n_bins)
    super().__init__(domain, codomain)

    self._n_bins = n_bins
    self._soft = soft
    self._temperature = temperature

    # compute bin centers and register as buffer
    assert domain.low is not None and domain.high is not None
    centers = torch.linspace(domain.low, domain.high, n_bins)
    self.register_buffer("centers", centers)

log_prob 

log_prob(x: Tensor, y: Tensor) -> Tensor

Log-probability of bin assignment y given input x.

PARAMETER DESCRIPTION
x

Continuous inputs. Shape (batch, 1).

TYPE: Tensor

y

Bin indices. Shape (batch,).

TYPE: Tensor

RETURNS DESCRIPTION
Tensor

Log-probabilities. Shape (batch,).
Source code in src/quivers/continuous/boundaries.py 
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
def log_prob(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
    """Log-probability of bin assignment y given input x.

    Parameters
    ----------
    x : torch.Tensor
        Continuous inputs. Shape (batch, 1).
    y : torch.Tensor
        Bin indices. Shape (batch,).

    Returns
    -------
    torch.Tensor
        Log-probabilities. Shape (batch,).
    """
    # soft assignment probabilities
    probs = self._soft_assign(x)  # (batch, n_bins)
    y_long = y.long()
    selected = probs[torch.arange(len(y_long)), y_long]
    return torch.log(selected.clamp(min=EPS))

rsample 

rsample(x: Tensor, sample_shape: Size = Size()) -> Tensor

Assign continuous inputs to discrete bins.

PARAMETER DESCRIPTION
x

Continuous inputs. Shape (batch, 1).

TYPE: Tensor

sample_shape

Ignored (assignment is deterministic).

TYPE: Size DEFAULT: Size()

RETURNS DESCRIPTION
Tensor

Bin indices. Shape (batch,) or (*sample_shape, batch).
Source code in src/quivers/continuous/boundaries.py 
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
def rsample(
    self,
    x: torch.Tensor,
    sample_shape: torch.Size = torch.Size(),
) -> torch.Tensor:
    """Assign continuous inputs to discrete bins.

    Parameters
    ----------
    x : torch.Tensor
        Continuous inputs. Shape (batch, 1).
    sample_shape : torch.Size
        Ignored (assignment is deterministic).

    Returns
    -------
    torch.Tensor
        Bin indices. Shape (batch,) or (*sample_shape, batch).
    """
    if self._soft:
        probs = self._soft_assign(x)
        bins = torch.multinomial(probs, 1).squeeze(-1)

    else:
        # hard assignment: closest bin center
        dists = (x - cast(torch.Tensor, self.centers).unsqueeze(0)).abs()
        bins = dists.argmin(dim=-1)

    if len(sample_shape) > 0:
        return bins.unsqueeze(0).expand(*sample_shape, -1)

    return bins

Embed 

Embed(domain: FinSet, codomain: Euclidean)

Bases: ContinuousMorphism

Map a finite set to a continuous space via kernel density placement.

Each element i of the domain FinSet is associated with a point in the continuous codomain. Sampling from Embed(x=i) produces a value near that point, with learnable spread.

Concretely, Embed places a Gaussian kernel at each bin center:

p(y | i) = Normal(center_i, sigma_i)

The centers and log-sigmas are learnable parameters.

PARAMETER DESCRIPTION
domain

Source discrete set.

TYPE: FinSet

codomain

Target continuous space (should be bounded for initialization).

TYPE: Euclidean

Source code in src/quivers/continuous/boundaries.py 
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
def __init__(
    self,
    domain: FinSet,
    codomain: Euclidean,
) -> None:
    super().__init__(domain, codomain)
    n = domain.size
    d = codomain.dim

    # initialize centers evenly spaced across codomain
    if codomain.is_bounded:
        assert codomain.low is not None and codomain.high is not None
        if d == 1:
            init_centers = torch.linspace(
                codomain.low,
                codomain.high,
                n,
            ).unsqueeze(-1)

        else:
            init_centers = (
                torch.rand(n, d) * (codomain.high - codomain.low) + codomain.low
            )

    else:
        init_centers = torch.randn(n, d) * 0.5

    self.centers = nn.Parameter(init_centers)
    self.log_sigma = nn.Parameter(torch.zeros(n, d))

log_prob 

log_prob(x: Tensor, y: Tensor) -> Tensor

Log-density of y under the kernel centered at x's embedding.

PARAMETER DESCRIPTION
x

Domain indices. Shape (batch,).

TYPE: Tensor

y

Continuous outputs. Shape (batch, d).

TYPE: Tensor

RETURNS DESCRIPTION
Tensor

Log-densities. Shape (batch,).
Source code in src/quivers/continuous/boundaries.py 
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
def log_prob(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
    """Log-density of y under the kernel centered at x's embedding.

    Parameters
    ----------
    x : torch.Tensor
        Domain indices. Shape (batch,).
    y : torch.Tensor
        Continuous outputs. Shape (batch, d).

    Returns
    -------
    torch.Tensor
        Log-densities. Shape (batch,).
    """
    mu = self.centers[x.long()]  # (batch, d)
    sigma = self.log_sigma[x.long()].exp().clamp(min=EPS)  # (batch, d)

    import math

    log_p = (
        -0.5 * ((y - mu) / sigma) ** 2 - sigma.log() - 0.5 * math.log(2 * math.pi)
    )

    return log_p.sum(dim=-1)

rsample 

rsample(x: Tensor, sample_shape: Size = Size()) -> Tensor

Sample from the Gaussian kernel at x's embedding point.

PARAMETER DESCRIPTION
x

Domain indices. Shape (batch,).

TYPE: Tensor

sample_shape

Additional sample dimensions.

TYPE: Size DEFAULT: Size()

RETURNS DESCRIPTION
Tensor

Continuous samples. Shape (*sample_shape, batch, d).
Source code in src/quivers/continuous/boundaries.py 
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
def rsample(
    self,
    x: torch.Tensor,
    sample_shape: torch.Size = torch.Size(),
) -> torch.Tensor:
    """Sample from the Gaussian kernel at x's embedding point.

    Parameters
    ----------
    x : torch.Tensor
        Domain indices. Shape (batch,).
    sample_shape : torch.Size
        Additional sample dimensions.

    Returns
    -------
    torch.Tensor
        Continuous samples. Shape (*sample_shape, batch, d).
    """
    mu = self.centers[x.long()]  # (batch, d)
    sigma = self.log_sigma[x.long()].exp().clamp(min=EPS)

    eps = torch.randn(
        *sample_shape,
        *mu.shape,
        device=mu.device,
        dtype=mu.dtype,
    )

    return mu + sigma * eps