Continuous Morphisms¶

@abstractmethod
def log_prob(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
    """Log-probability (density) of y given x.

    Parameters
    ----------
    x : torch.Tensor
        Inputs. Shape (batch,) for discrete domain or
        (batch, domain_dim) for continuous domain.
    y : torch.Tensor
        Outputs. Shape (batch,) for discrete codomain or
        (batch, codomain_dim) for continuous codomain.

    Returns
    -------
    torch.Tensor
        Log-probabilities/densities. Shape (batch,).
    """
    ...

@abstractmethod
def rsample(
    self, x: torch.Tensor, sample_shape: torch.Size = torch.Size()
) -> torch.Tensor:
    """Reparameterized samples from p(. | x).

    Gradients flow through the returned samples back to the
    parameters of this morphism (and to x if the domain is
    continuous).

    Parameters
    ----------
    x : torch.Tensor
        Inputs. Shape (batch,) or (batch, domain_dim).
    sample_shape : torch.Size
        Additional leading sample dimensions.

    Returns
    -------
    torch.Tensor
        Samples. Shape (*sample_shape, batch, codomain_dim) for
        continuous codomain, or (*sample_shape, batch) for discrete.
    """
    ...

def sample(
    self, x: torch.Tensor, sample_shape: torch.Size = torch.Size()
) -> torch.Tensor:
    """Non-reparameterized samples (no gradient through samples).

    Parameters
    ----------
    x : torch.Tensor
        Inputs.
    sample_shape : torch.Size
        Additional leading sample dimensions.

    Returns
    -------
    torch.Tensor
        Samples (detached from computation graph).
    """
    with torch.no_grad():
        return self.rsample(x, sample_shape)

def __rshift__(self, other: object) -> ContinuousMorphism:
    """Composition via ancestral sampling: self >> other."""
    if isinstance(other, ContinuousMorphism):
        return SampledComposition(self, other)
    from quivers.core.morphisms import Morphism

    if isinstance(other, Morphism):
        return SampledComposition(self, DiscreteAsContinuous(other))
    return NotImplemented

def __rrshift__(self, other: object) -> ContinuousMorphism:
    """Handle discrete_morphism >> continuous_morphism."""
    from quivers.core.morphisms import Morphism

    if isinstance(other, Morphism):
        return SampledComposition(DiscreteAsContinuous(other), self)
    return NotImplemented

def __matmul__(self, other: object) -> ProductContinuousMorphism:
    """Independent product: self @ other."""
    if isinstance(other, ContinuousMorphism):
        return ProductContinuousMorphism(self, other)
    from quivers.core.morphisms import Morphism

    if isinstance(other, Morphism):
        return ProductContinuousMorphism(self, DiscreteAsContinuous(other))
    return NotImplemented

def rsample(
    self, x: torch.Tensor, sample_shape: torch.Size = torch.Size()
) -> torch.Tensor:
    """Ancestral sampling: y ~ f(x, .), then z ~ g(y, .).

    Parameters
    ----------
    x : torch.Tensor
        Inputs to the composition.
    sample_shape : torch.Size
        Additional sample dimensions.

    Returns
    -------
    torch.Tensor
        Samples from the composed morphism.
    """
    y = self.left.rsample(x, sample_shape)
    if len(sample_shape) > 0:
        leading = y.shape[: len(sample_shape)]
        batch = x.shape[0]
        flat_size = int(torch.tensor(leading).prod().item()) * batch
        if y.dim() > len(sample_shape) + 1:
            event_dims = y.shape[len(sample_shape) + 1 :]
            flat_y = y.reshape(flat_size, *event_dims)
        else:
            flat_y = y.reshape(flat_size)
    else:
        flat_y = y
    z = self.right.rsample(flat_y)
    if len(sample_shape) > 0:
        batch = x.shape[0]
        if z.dim() > 1:
            event_dims = z.shape[1:]
            z = z.reshape(*sample_shape, batch, *event_dims)
        else:
            z = z.reshape(*sample_shape, batch)
    return z

def log_prob(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
    """Log-probability of y given x through the composition.

    When the intermediate space is discrete, computes the exact
    marginalization. When continuous, uses Monte Carlo estimation.

    Parameters
    ----------
    x : torch.Tensor
        Inputs. Shape (batch,) or (batch, dom_dim).
    y : torch.Tensor
        Outputs. Shape (batch,) or (batch, cod_dim).

    Returns
    -------
    torch.Tensor
        Log-probabilities. Shape (batch,).
    """
    intermediate = self.left.codomain
    if isinstance(intermediate, SetObject):
        return self._log_prob_exact(x, y, intermediate)
    else:
        return self._log_prob_mc(x, y)

def rsample(
    self, x: torch.Tensor, sample_shape: torch.Size = torch.Size()
) -> torch.Tensor:
    """Sample from all components and concatenate outputs.

    Parameters
    ----------
    x : torch.Tensor
        Input tensor (broadcast to all components).
    sample_shape : torch.Size
        Additional leading sample dimensions.

    Returns
    -------
    torch.Tensor
        Concatenated outputs from all components.
    """
    outs = []
    for comp in self._components:
        y = cast(ContinuousMorphism, comp).rsample(x, sample_shape)
        if y.dim() == 1:
            y = y.unsqueeze(-1)
        outs.append(y)
    return torch.cat(outs, dim=-1)

def log_prob(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
    """Log-probability: sum of component log-probs.

    Parameters
    ----------
    x : torch.Tensor
        Input (same for all components).
    y : torch.Tensor
        Concatenated output values.

    Returns
    -------
    torch.Tensor
        Sum of log-probabilities. Shape ``(batch,)``.
    """
    lp = torch.zeros(x.shape[0], device=x.device)
    offset = 0
    for comp_mod, d in zip(self._components, self._cod_dims):
        comp = cast(ContinuousMorphism, comp_mod)
        y_slice = y[..., offset : offset + d]
        if _is_discrete(comp.codomain):
            y_slice = y_slice.squeeze(-1).long()
        lp = lp + comp.log_prob(x, y_slice)
        offset += d
    return lp

def log_prob(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
    """Log-probability from the discrete tensor.

    Parameters
    ----------
    x : torch.Tensor
        Domain indices. Shape (batch,).
    y : torch.Tensor
        Codomain indices. Shape (batch,).

    Returns
    -------
    torch.Tensor
        Log-probabilities. Shape (batch,).
    """
    t = self._inner.tensor
    probs = t[x.long(), y.long()]
    return torch.log(probs.clamp(min=1e-07))

def rsample(
    self, x: torch.Tensor, sample_shape: torch.Size = torch.Size()
) -> torch.Tensor:
    """Sample from the categorical distribution defined by the tensor.

    Note: not reparameterizable. Gradients do not flow through
    the returned samples.

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

    Returns
    -------
    torch.Tensor
        Sampled codomain indices. Shape (*sample_shape, batch).
    """
    t = self._inner.tensor
    probs = t[x.long()]
    n_samples = (
        int(torch.Size(sample_shape).numel()) if len(sample_shape) > 0 else 1
    )
    samples = torch.multinomial(probs, n_samples, replacement=True)
    if len(sample_shape) == 0:
        return samples.squeeze(-1)
    else:
        return samples.T.reshape(*sample_shape, -1)