Continuous Programs¶

def rsample(  # type: ignore[override]
    self,
    x: torch.Tensor,
    sample_shape: torch.Size = torch.Size(),
    observations: dict[str, torch.Tensor] | None = None,
) -> torch.Tensor | dict[str, torch.Tensor]:
    """Run the program forward, returning the designated output(s).

    Each draw step is executed in order. Steps that reference
    the program input use ``x`` directly; steps that reference
    bound variables use those variables' sampled values.

    Parameters
    ----------
    x : torch.Tensor
        Program input.
    sample_shape : torch.Size
        Additional leading sample dimensions (applied to the
        first draw only; subsequent draws inherit the shape).
    observations : dict[str, torch.Tensor] or None
        Values to clamp observed variables to. Keys are variable
        names, values are tensors of the appropriate shape.

    Returns
    -------
    torch.Tensor or dict[str, torch.Tensor]
        The value of the return variable(s). Returns a tensor
        for single-variable returns, or a dict keyed by variable
        name for tuple returns.
    """
    if observations is None:
        observations = {}

    env: dict[str, torch.Tensor] = {}

    # pre-populate env with named params (split product input)
    if self._params is not None and self._param_dims is not None:
        splits = torch.split(x, self._param_dims, dim=-1)

        assert self._param_is_continuous is not None
        for pname, chunk, is_cont in zip(
            self._params, splits, self._param_is_continuous
        ):
            # only squeeze discrete components (continuous dim=1 should stay 2D)
            if not is_cont and chunk.shape[-1] == 1:
                env[pname] = chunk.squeeze(-1)

            else:
                env[pname] = chunk

    for i, spec in enumerate(self._step_specs):
        if isinstance(spec, _LetSpec):
            # deterministic binding: constant, alias, or expression
            if isinstance(spec.value, str):
                env[spec.var] = env[spec.value]

            elif callable(spec.value):
                env[spec.var] = cast(torch.Tensor, spec.value(env))

            else:
                env[spec.var] = torch.full(
                    (x.shape[0],),
                    spec.value,
                    device=x.device,
                )

            continue

        assert self._modules[spec.morphism_name] is not None
        morph = cast(ContinuousMorphism, self._modules[spec.morphism_name])
        inp = self._resolve_input(spec, x, env)

        # check if any vars in this step are observed
        if len(spec.vars) == 1:
            var_name = spec.vars[0]

            if spec.is_observed and var_name in observations:
                # clamp to observed value
                env[var_name] = observations[var_name]
                continue

        else:
            # destructuring: if observed vars are present, clamp them
            any_clamped = False

            for v in spec.vars:
                if spec.is_observed and v in observations:
                    env[v] = observations[v]
                    any_clamped = True

            if any_clamped:
                # for partially observed destructuring, sample the rest
                all_clamped = all(v in observations for v in spec.vars)

                if not all_clamped:
                    result = morph.rsample(inp)
                    # only bind un-clamped vars
                    if isinstance(result, dict):
                        result_dict = cast(dict[str, torch.Tensor], result)
                        for v in spec.vars:
                            if v not in observations:
                                env[v] = result_dict[v]

                    else:
                        dims = self._compute_component_dims(morph.codomain)
                        splits = torch.split(result, dims, dim=-1)

                        for v, chunk in zip(spec.vars, splits):
                            if v not in observations:
                                env[v] = (
                                    chunk.squeeze(-1)
                                    if chunk.shape[-1] == 1
                                    else chunk
                                )

                continue

        # only apply sample_shape to the first draw from input
        if i == 0 and spec.args is None and len(sample_shape) > 0:
            result = morph.rsample(inp, sample_shape)

        else:
            result = morph.rsample(inp)

        self._bind_result(spec, result, env)

    # return
    if self._return_is_single:
        return env[self._return_vars[0]]

    # use labels as keys if available, otherwise variable names
    keys = self._return_labels if self._return_labels else self._return_vars
    return {k: env[v] for k, v in zip(keys, self._return_vars)}

def log_prob(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
    """Log-probability is not supported for monadic programs.

    Computing log p(y | x) for a monadic program requires
    marginalizing over all intermediate variables, which is
    intractable in general. Use ``rsample`` for forward sampling
    and condition via score function estimators or variational
    methods.

    Raises
    ------
    NotImplementedError
        Always.
    """
    raise NotImplementedError(
        "log_prob is not supported for monadic programs; "
        "computing p(y | x) requires marginalizing over all "
        "intermediate draws, which is intractable in general. "
        "use rsample() for forward sampling."
    )

def log_joint(
    self,
    x: torch.Tensor,
    intermediates: dict[str, torch.Tensor],
) -> torch.Tensor:
    """Joint log-density given all intermediate values.

    When all intermediate variables are observed (e.g. during
    inference with HMC/NUTS), computes the joint log-density:

        log p(x_1, ..., x_n | input) = sum_i log p(x_i | pa(x_i))

    where pa(x_i) is the parent variable of step i (either the
    program input or a previously drawn variable).

    For destructuring draw steps (tuple-returning sub-programs),
    the intermediates dict should contain entries for each
    individual variable name.

    Parameters
    ----------
    x : torch.Tensor
        Program input.
    intermediates : dict[str, torch.Tensor]
        Values for ALL bound variables (keyed by variable name
        or by return label if labels are set).

    Returns
    -------
    torch.Tensor
        Joint log-density. Shape (batch,).
    """
    total = torch.zeros(x.shape[0], device=x.device)

    # if labels are used, map label keys back to variable names
    env = dict(intermediates)

    if self._return_labels:
        for label, var in zip(self._return_labels, self._return_vars):
            if label in env and var not in env:
                env[var] = env[label]

    if self._params is not None and self._param_dims is not None:
        splits = torch.split(x, self._param_dims, dim=-1)

        assert self._param_is_continuous is not None
        for pname, chunk, is_cont in zip(
            self._params, splits, self._param_is_continuous
        ):
            if pname not in env:
                if not is_cont and chunk.shape[-1] == 1:
                    env[pname] = chunk.squeeze(-1)

                else:
                    env[pname] = chunk

    for spec in self._step_specs:
        if isinstance(spec, _LetSpec):
            # deterministic binding: populate env, contribute 0
            if spec.var not in env:
                if isinstance(spec.value, str):
                    env[spec.var] = env[spec.value]

                elif callable(spec.value):
                    env[spec.var] = cast(torch.Tensor, spec.value(env))

                else:
                    env[spec.var] = torch.full(
                        (x.shape[0],),
                        spec.value,
                        device=x.device,
                    )

            continue

        assert self._modules[spec.morphism_name] is not None
        morph = cast(ContinuousMorphism, self._modules[spec.morphism_name])
        inp = self._resolve_input(spec, x, env)

        if len(spec.vars) == 1:
            val = env[spec.vars[0]]
            total = total + morph.log_prob(inp, val)

        else:
            # destructuring step: if sub-program, call its log_joint
            # with the individual intermediate values
            if hasattr(morph, "log_joint") and hasattr(morph, "_return_vars"):
                # reconstruct the sub-program's intermediates from
                # the overall intermediates dict
                sub_morph = cast(MonadicProgram, morph)
                sub_intermediates = {}

                for sub_spec in sub_morph._step_specs:
                    if isinstance(sub_spec, _LetSpec):
                        continue

                    for sv in sub_spec.vars:
                        if sv in env:
                            sub_intermediates[sv] = env[sv]

                total = total + sub_morph.log_joint(inp, sub_intermediates)

            else:
                # product-codomain morphism: reconstruct stacked output
                parts = [env[v] for v in spec.vars]
                val = self._stack_tensors(parts)
                total = total + morph.log_prob(inp, val)

    return total