Core¶

def trace(
    fun: Callable,
    traced_vars=None,
    *,
    sites: bool = False,
    kind_labels: str = "probabilistic",
):
    base_fun = _get_base_fun(fun)
    get_jaxpr = _cached_jaxpr_getter(base_fun)
    interventions, observations, replay = _collect_stochastic_maps(fun)

    @wraps(fun)
    def wrapped(*args, **kwargs):
        processing_rule = TraceProcessingRule(
            traced_vars=traced_vars,
            sites=sites,
            interventions=interventions,
            observations=observations,
            replay=replay,
            kind_labels=kind_labels,
        )
        with enable_rv_tracing():
            jaxpr = get_jaxpr(*args, **kwargs)
        trace_result = cast(
            tuple[list, dict],
            interpret(
                jaxpr.jaxpr,
                jaxpr.consts,
                jaxpr.jaxpr.invars,
                _flatten_call_inputs(args, kwargs),
                jaxpr.jaxpr.outvars,
                process_eqn=processing_rule,
                reducer=trace_state_reducer,
                initial_state={},
                return_state=True,
            ),
        )
        traced_samples = trace_result[1]

        return traced_samples

    return wrapped

def joint_sample(fun: Callable, rvs: Optional[Iterable] = None) -> Callable:
    """Samples all random variables called in the probabilistic function. If rvs is
    given, it only samples the random variables in rvs.

    Args:
        fun (Callable): Probabilistic function
        rvs (Optional[Iterable], optional): Subset of random variables in the
            probabilistic program. Defaults to None.

    Returns:
        Callable: Sampling function
    """
    base_fun = _get_base_fun(fun)
    get_jaxpr = _cached_jaxpr_getter(base_fun)
    interventions, observations, replay = _collect_stochastic_maps(fun)
    fixed_values = {}
    fixed_values.update(replay)
    fixed_values.update(observations)
    fixed_values.update(interventions)
    legacy_fixed = getattr(fun, "_probjax_interventions", None)

    def wrapped(*args, **kwargs):
        processing_rule = JointSampleProcessingRule(
            rvs=rvs,
            fixed_values=fixed_values,
            fixed_names=legacy_fixed,
        )
        with enable_rv_tracing():
            jaxpr = get_jaxpr(*args, **kwargs)
        joint_result = cast(
            tuple[list, dict],
            interpret(
                jaxpr.jaxpr,
                jaxpr.consts,
                jaxpr.jaxpr.invars,
                _flatten_call_inputs(args, kwargs),
                jaxpr.jaxpr.outvars,
                process_eqn=processing_rule,
                reducer=joint_sample_state_reducer,
                initial_state={},
                return_state=True,
            ),
        )
        joint_samples = joint_result[1]

        return joint_samples

    return wrapped

def log_joint_fn(
    fun: Callable,
    *args,
    strict: bool = True,
    allow_partial: bool = False,
    **kwargs,
):
    """Compute the model log-joint (up to a constant)."""
    return log_potential_fn(
        fun,
        *args,
        strict=strict,
        allow_partial=allow_partial,
        **kwargs,
    )

def log_potential_fn(
    fun: Callable,
    *args,
    strict: bool = True,
    allow_partial: bool = False,
    **kwargs,
):
    """Compute the unnormalized log density of a probabilistic function.

    This is the legacy name for :func:`log_joint_fn`.

    This does not include the normalizing constant.

    Args:
        fun (Callable): Probabilistic function

    Returns:
        Callable: Log potential function
    """
    base_fun = _get_base_fun(fun)
    interventions, observations, replay = _collect_stochastic_maps(fun)
    legacy_interventions = getattr(fun, "_probjax_interventions", None)

    with enable_rv_tracing():
        jaxpr = jax.make_jaxpr(base_fun)(jax.random.PRNGKey(0), *args, **kwargs)
    model_inputs = _flatten_call_inputs((jax.random.PRNGKey(0),) + args, kwargs)

    def log_potential(**joint_samples):
        intervention_config = interventions if interventions else legacy_interventions
        processing_rule = LogPotentialProcessingRule(
            joint_samples=joint_samples,
            interventions=intervention_config,
            observations=observations,
            replay=replay,
            strict=strict,
            allow_partial=allow_partial,
        )

        log_potential_result = cast(
            tuple[list, jax.Array],
            interpret(
                jaxpr.jaxpr,
                jaxpr.consts,
                jaxpr.jaxpr.invars,
                model_inputs,
                jaxpr.jaxpr.outvars,
                process_eqn=processing_rule,
                reducer=log_potential_state_reducer,
                initial_state=jnp.asarray(0.0),
                return_state=True,
            ),
        )
        log_prob = log_potential_result[1]

        return jnp.nan_to_num(log_prob, nan=-jnp.inf, posinf=jnp.inf, neginf=-jnp.inf)

    return log_potential

def condition(fun: Callable, observations: Mapping[str, Array]) -> Callable:
    """Condition a probabilistic program on observed site values.

    The preferred probabilistic name is :func:`observe`.
    """
    return _apply_stochastic_substitution(fun, update_observations=observations)

def observe(fun: Callable, observations: Mapping[str, Array]) -> Callable:
    """Observe stochastic sites in a probabilistic program.

    This is a probabilistic alias for :func:`condition`.
    """
    return condition(fun, observations)

def intervene(fun: Callable, rvs: Mapping[str, Array]) -> Callable:
    """Fix stochastic sites via intervention values.

    This is equivalent to a causal ``do`` operation. The name ``do`` is the
    probabilistic alias for this function.

    This does not sample the random variables, but fixes them to the given values.

    The wrapped function uses interpreter-level overrides for the selected
    random variables while leaving all other equations unchanged.

    Args:
        fun (Callable): A function to transform.
        rvs (dict[str, Array]): A dictionary of random variable names and values to
            intervene.

    Returns:
        Callable: Wrapped probabilistic function with interventions.
    """
    return _apply_stochastic_substitution(fun, update_interventions=rvs)

def do(fun: Callable, interventions: Mapping[str, Array]) -> Callable:
    """Apply a causal ``do`` intervention to stochastic sites.

    This is a probabilistic alias for :func:`intervene`.
    """
    return intervene(fun, interventions)

def substitute(
    fun: Callable, values: Mapping[str, Array], mode: str = "condition"
) -> Callable:
    """Substitute stochastic sites with fixed values.

    Args:
        fun: Probabilistic function.
        values: Site-value mapping.
        mode: `"condition"` (or legacy `"replay"`) includes substituted sites
            in log-probability; `"do"` (or legacy `"intervene"`) applies
            intervention semantics and drops substituted-site log terms.
    """
    normalized_mode = _normalize_substitute_mode(mode)

    if normalized_mode == "condition":
        return _apply_stochastic_substitution(fun, update_replay=values)

    if normalized_mode == "do":
        return intervene(fun, values)

    raise AssertionError("unreachable")

def scope(name: str):
    """Create a naming scope for stochastic sites."""
    return name_stack.scope(name)

def inverse(fun: Callable, static_argnums=(), invertible_arg=None, input_template=None):
    """Return a function computing the inverse of ``fun``.

    Traces ``fun`` to a jaxpr and walks it backwards, replacing each primitive
    with its registered inverse rule, so the result is ordinary JAX code with no
    interpreter left at runtime:

    >>> inverse(lambda x: 2 * jnp.exp(x))(jnp.asarray(4.0))
    Array(0.6931472, dtype=float32)

    Args:
        fun: The function to invert. May itself be a ``custom_inverse``, in
            which case its registered inverse is used directly.
        static_argnums: Positional arguments held fixed rather than inverted.
        invertible_arg: Which positional argument to solve for (default 0).
            A tuple of indices solves several arguments jointly -- e.g.
            ``invertible_arg=(0, 1)`` inverts ``(x, y) -> (x + y, x - y)``
            back to ``(x, y)`` -- and returns them as a tuple in order.
        input_template: Optional tracing example. It must preserve boundary shapes;
            templates cannot enable expansion or other shape-changing inverses.

    Returns:
        A callable mapping outputs back to the invertible argument.

    Raises:
        ValueError: if traced input/output shapes or pytree structures differ.
            Automatic inversion uses the supplied output as the presumed input
            signature; it requires a shape-preserving function. Some violations
            (such as a broadcast that becomes an identity at that signature)
            cannot be detected without the original input specification.

    Note:
        Detectable boundary shape/tree mismatches raise ValueError.
        Elementwise affine sections use symbolic coefficients, without Jacobians.
        Affine sections compose with nonlinear inverse rules; analysis and
        schedules are cached, including normalized nested jit programs.

        **An unresolved inversion returns NaN.** The interpreter works
        one equation at a time, so it inverts a *tree* of operations; a value
        used twice stalls it, because the bivariate rules need exactly one
        unknown operand.

        On a stall, structurally proven affine sections are recovered by linear
        solves and propagation resumes. This supports ``exp(3*x-x)`` and
        ``3*exp(x)-exp(x)``, including sequential compositions and nested jit.
        Pointwise maps invert elementwise in O(n); small coupled maps build
        the matrix with one vmapped sweep; large coupled maps solve
        matrix-free. When a stall cannot be sectioned, a whole-program affine
        solve replaces NaN with a value. Analysis and scheduling decisions
        are cached; generated inverses are ordinary JAX computations and can
        be jitted, vmapped, and differentiated.


        These remain silently unsupported and produce NaN:

        * a value used more than once **nonlinearly** -- ``x * x``, or a
          residual ``x + f(x)``. A residual is invertible by fixed-point
          iteration when its branch is a contraction, but a jaxpr carries no
          Lipschitz bound, so register it with :class:`custom_inverse` instead.
        * ``lax.fori_loop``, and any ``lax.scan`` carrying something that is not
          itself invertible (a counter, a running sum). Plain ``scan`` and
          ``lax.cond`` do work.
        * ``inverse(inverse(f))``.

        ``lax.while_loop`` raises rather than returning NaN. Checking
        ``jnp.isfinite`` on the result is the reliable way to detect the rest.
    """
    if input_template is not None and isinstance(fun, custom_inverse):
        raise ValueError("input_template has no effect for custom_inverse inputs")
    maybe_custom = maybe_inverse_custom_inverse(
        fun,
        static_argnums=static_argnums,
        invertible_arg=invertible_arg,
    )
    if maybe_custom is not None:
        return maybe_custom

    get_jaxpr = _cached_jaxpr_getter(
        fun, static_argnums=static_argnums, return_shape=True
    )
    recovery_cache = {}
    schedule_cache = {}

    @wraps(fun)
    def wrapped(*args, **kwargs):
        processing_rule = InverseProcessingRule()
        (
            jaxpr,
            known_invars,
            target_invars,
            args_for_propagate,
            target_trees,
        ) = _trace_inverse_program(
            get_jaxpr, args, kwargs, invertible_arg, static_argnums, input_template
        )
        graph = cached_inverse_graph(jaxpr, target_invars, recovery_cache)
        if input_template is not None and not _outputs_match_program(
            args_for_propagate[len(known_invars) :], jaxpr.jaxpr.outvars
        ):
            return _unflatten_targets(
                target_trees, _unresolvable_targets(target_invars)
            )
        out, env = cast(
            tuple[list, Any],
            propagate(
                graph.jaxpr,
                graph.consts,
                known_invars + graph.jaxpr.outvars,
                args_for_propagate,
                target_invars,
                process_eqn=processing_rule,
                schedule_cache=schedule_cache,
                stall_recovery=make_affine_recovery(
                    graph,
                    known_invars,
                    args_for_propagate,
                    target_invars,
                    processing_rule,
                    recovery_cache,
                    schedule_cache=schedule_cache,
                    with_logdet=isinstance(
                        processing_rule, InverseAndLogAbsDetProcessingRule
                    ),
                ),
                cost_fn=inverse_cost_fn,
                process_all_eqns=True,
                return_env=True,
            ),
        )
        out, complete = _materialize_inverse_targets(out, target_invars, env)
        if not complete:
            solved = _affine_fallback(
                jaxpr,
                known_invars,
                args_for_propagate,
                target_invars,
                need_logdet=False,
            )
            if solved is not None:
                out = solved[0]

        return _unflatten_targets(target_trees, out)

    return wrapped

def inverse_and_logabsdet(
    fun: Callable, static_argnums=(), invertible_arg=None, input_template=None
):
    """Return a function computing the inverse of ``fun`` and its log-det.

    The log-determinant is that of the **inverse** map -- ``log|d(inv)/dy|``,
    summed over the event -- which is the term a change of variables needs:

    >>> inverse_and_logabsdet(lambda x: 2 * jnp.exp(x))(jnp.asarray(4.0))
    (Array(0.6931472, dtype=float32), Array(-1.3862944, dtype=float32))

    Args:
        fun: The function to invert, possibly a ``custom_inverse``.
        static_argnums: Positional arguments held fixed rather than inverted.
        invertible_arg: Which positional argument to solve for (default 0).
            A tuple of indices solves several arguments jointly and returns
            them as a tuple in order.
        input_template: Optional tracing example. It must preserve boundary shapes;
            templates cannot enable expansion or other shape-changing inverses.

    Returns:
        ``(inverse_value, log_abs_det)``. Both are NaN if the inversion could
        not be completed.

    Raises:
        ValueError: if traced boundary shapes or pytree structures differ,
            with the same input-signature limitation as :func:`inverse`.
        NotImplementedError: if a primitive on the inverse path has no
            log-determinant rule and is not elementwise. Guessing one by
            differentiating the inverse elementwise -- the old behaviour --
            silently returned a number that was not a log-determinant.

    Note:
        Everything :func:`inverse` cannot do applies here too, and the log-det
        additionally requires a rule for every primitive involved. A primitive
        that inverts fine may still have no log-det: ``dynamic_slice`` recovers
        only its window, leaving the input partially known and no square
        Jacobian to take a determinant of.

    Note:
        When the program is structurally volume-preserving in the target --
        rearrangements, translations, ``neg``, ``±1`` scalings -- the log-det
        is proven zero from the jaxpr and no accumulation is staged at all, so
        the compiled inverse matches a hand-written one equation for equation.
    """
    if input_template is not None and isinstance(fun, custom_inverse):
        raise ValueError("input_template has no effect for custom_inverse inputs")
    maybe_custom = maybe_inverse_custom_inverse(
        fun,
        static_argnums=static_argnums,
        invertible_arg=invertible_arg,
    )
    if maybe_custom is not None:

        @wraps(fun)
        def custom_wrapped(*args, **kwargs):
            value, logdet = fun.inv_and_logdet(*args, **kwargs)
            total_logdet = sum(
                (jnp.sum(leaf) for leaf in jax.tree_util.tree_leaves(logdet)),
                jnp.asarray(0.0),
            )
            return value, total_logdet

        return custom_wrapped

    get_jaxpr = _cached_jaxpr_getter(
        fun, static_argnums=static_argnums, return_shape=True
    )
    recovery_cache = {}
    schedule_cache = {}

    @wraps(fun)
    def wrapped(*args, **kwargs):
        processing_rule = InverseAndLogAbsDetProcessingRule(
            state_namespace=INVERSE_AND_LOGABSDET_STATE_NAMESPACE
        )
        (
            jaxpr,
            known_invars,
            target_invars,
            args_for_propagate,
            target_trees,
        ) = _trace_inverse_program(
            get_jaxpr, args, kwargs, invertible_arg, static_argnums, input_template
        )
        graph = cached_inverse_graph(jaxpr, target_invars, recovery_cache)
        invars = known_invars + graph.jaxpr.outvars
        outvars = target_invars

        if input_template is not None and not _outputs_match_program(
            args_for_propagate[len(known_invars) :], jaxpr.jaxpr.outvars
        ):
            return _unflatten_targets(
                target_trees, _unresolvable_targets(target_invars)
            ), jnp.asarray(jnp.nan)

        if is_volume_preserving(jaxpr.jaxpr, target_invars):
            # Proven |det J| = 1 from the jaxpr structure, so the log-det is
            # exactly zero and the accumulation machinery below would only
            # stage per-equation `+ 0.0`s. Run the plain inverse propagation
            # instead; on incompleteness fall through to the full path so
            # failure semantics (NaN, affine fallback) are unchanged.
            out, env = cast(
                tuple[list, Any],
                propagate(
                    jaxpr.jaxpr,
                    jaxpr.consts,
                    known_invars + jaxpr.jaxpr.outvars,
                    args_for_propagate,
                    target_invars,
                    process_eqn=InverseProcessingRule(),
                    cost_fn=inverse_cost_fn,
                    process_all_eqns=True,
                    return_env=True,
                ),
            )
            out, complete = _materialize_inverse_targets(out, target_invars, env)
            if complete:
                return _unflatten_targets(target_trees, out), jnp.asarray(0.0)

        inverse_result = cast(
            tuple[list, dict, Any],
            propagate(
                graph.jaxpr,
                graph.consts,
                invars,
                args_for_propagate,
                outvars,
                process_eqn=processing_rule,
                schedule_cache=schedule_cache,
                stall_recovery=make_affine_recovery(
                    graph,
                    known_invars,
                    args_for_propagate,
                    target_invars,
                    processing_rule,
                    recovery_cache,
                    schedule_cache=schedule_cache,
                    with_logdet=isinstance(
                        processing_rule, InverseAndLogAbsDetProcessingRule
                    ),
                ),
                cost_fn=inverse_cost_fn,
                process_all_eqns=True,
                reducer=inverse_and_logabsdet_state_reducer,
                initial_state={},
                return_state=True,
                return_env=True,
                state_namespace=INVERSE_AND_LOGABSDET_STATE_NAMESPACE,
            ),
        )
        out, log_dets, env = inverse_result
        out, complete = _materialize_inverse_targets(out, outvars, env)

        log_det = _sum_log_dets_for_vars(log_dets, outvars)
        if not complete:
            solved = _affine_fallback(jaxpr, known_invars, args_for_propagate, outvars)
            if solved is not None:
                out, log_det = solved
            else:
                log_det = jnp.asarray(jnp.nan)
        return _unflatten_targets(target_trees, out), log_det

    return wrapped