Program¶

def rsample(
    self,
    x: torch.Tensor,
    sample_shape: torch.Size = torch.Size(),
) -> torch.Tensor:
    """Reparameterized sample (continuous programs only).

    Parameters
    ----------
    x : torch.Tensor
        Input tensor.
    sample_shape : torch.Size
        Extra leading sample dimensions.

    Returns
    -------
    torch.Tensor
        Sampled output.

    Raises
    ------
    TypeError
        If the underlying morphism is not continuous.
    """
    if not isinstance(self._morphism, ContinuousMorphism):
        raise TypeError("rsample is only available for continuous programs")

    return cast(torch.Tensor, self._morphism.rsample(x, sample_shape))

def log_prob(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
    """Log-probability (continuous programs only).

    Parameters
    ----------
    x : torch.Tensor
        Input tensor.
    y : torch.Tensor
        Output tensor.

    Returns
    -------
    torch.Tensor
        Log-probability of y given x.

    Raises
    ------
    TypeError
        If the underlying morphism is not continuous.
    """
    if not isinstance(self._morphism, ContinuousMorphism):
        raise TypeError("log_prob is only available for continuous programs")

    return cast(torch.Tensor, self._morphism.log_prob(x, y))

def forward(self, n_steps: int | None = None) -> torch.Tensor:
    """Materialize the composed tensor (discrete programs only).

    Parameters
    ----------
    n_steps : int or None
        If provided, sets the step count on all
        ``RepeatMorphism`` instances in the morphism tree
        before materializing. This allows runtime-variable
        sequence lengths for models using ``repeat(f)``.

    Returns
    -------
    torch.Tensor
        Tensor of shape ``(*domain.shape, *codomain.shape)`` with
        values in [0, 1].

    Raises
    ------
    TypeError
        If the underlying morphism is continuous.
    """
    if self._is_continuous:
        raise TypeError(
            "forward() is not supported for continuous programs; "
            "use rsample() or log_prob() instead"
        )

    if n_steps is not None:
        self._set_repeat_steps(n_steps)

    return cast(Morphism, self._morphism).tensor

def log_membership(self) -> torch.Tensor:
    """Log of the membership tensor.

    Returns
    -------
    torch.Tensor
        Log-membership values.  Entries near 0 map to large
        negative values.
    """
    t = cast(Morphism, self._morphism).tensor
    return torch.log(t.clamp(min=1e-7))

def nll_loss(
    self,
    domain_indices: torch.Tensor,
    codomain_indices: torch.Tensor,
) -> torch.Tensor:
    """Negative log-likelihood loss for observed (domain, codomain) pairs.

    For each pair (x, y), computes -log(tensor[x, y]).
    Suitable when the morphism represents fuzzy membership and
    observed pairs should have high membership.

    Parameters
    ----------
    domain_indices : torch.Tensor
        Integer tensor of shape (batch,) or (batch, n_domain_dims)
        indexing into the domain.
    codomain_indices : torch.Tensor
        Integer tensor of shape (batch,) or (batch, n_codomain_dims)
        indexing into the codomain.

    Returns
    -------
    torch.Tensor
        Scalar mean negative log-likelihood.
    """
    t = cast(Morphism, self._morphism).tensor

    # handle both flat and multi-dimensional indexing
    if domain_indices.ndim == 1:
        domain_indices = domain_indices.unsqueeze(-1)

    if codomain_indices.ndim == 1:
        codomain_indices = codomain_indices.unsqueeze(-1)

    # combine indices
    indices = torch.cat([domain_indices, codomain_indices], dim=-1)

    # index into the tensor
    idx_tuple = tuple(indices[:, i] for i in range(indices.shape[1]))
    values = t[idx_tuple]

    return -torch.log(values.clamp(min=1e-7)).mean()

def bce_loss(self, target: torch.Tensor) -> torch.Tensor:
    """Binary cross-entropy between the morphism tensor and a target.

    Parameters
    ----------
    target : torch.Tensor
        Target tensor of the same shape as the morphism tensor,
        with values in [0, 1].

    Returns
    -------
    torch.Tensor
        Scalar BCE loss.
    """
    t = cast(Morphism, self._morphism).tensor
    return F.binary_cross_entropy(t, target)