pymc-extras 0.2.7__py3-none-any.whl → 0.4.0__py3-none-any.whl

pymc_extras/inference/laplace_approx/laplace.py

@@ -0,0 +1,453 @@
+#   Copyright 2024 The PyMC Developers
+#
+#   Licensed under the Apache License, Version 2.0 (the "License");
+#   you may not use this file except in compliance with the License.
+#   You may obtain a copy of the License at
+#
+#       http://www.apache.org/licenses/LICENSE-2.0
+#
+#   Unless required by applicable law or agreed to in writing, software
+#   distributed under the License is distributed on an "AS IS" BASIS,
+#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#   See the License for the specific language governing permissions and
+#   limitations under the License.
+
+
+import logging
+
+from collections.abc import Callable
+from functools import partial
+from typing import Literal
+from typing import cast as type_cast
+
+import arviz as az
+import numpy as np
+import pymc as pm
+import pytensor
+import pytensor.tensor as pt
+import xarray as xr
+
+from better_optimize.constants import minimize_method
+from numpy.typing import ArrayLike
+from pymc.blocking import DictToArrayBijection
+from pymc.model.transform.optimization import freeze_dims_and_data
+from pymc.pytensorf import join_nonshared_inputs
+from pymc.util import get_default_varnames
+from pytensor.graph import vectorize_graph
+from pytensor.tensor import TensorVariable
+from pytensor.tensor.optimize import minimize
+from pytensor.tensor.type import Variable
+
+from pymc_extras.inference.laplace_approx.find_map import (
+    _compute_inverse_hessian,
+    _make_initial_point,
+    find_MAP,
+)
+from pymc_extras.inference.laplace_approx.scipy_interface import (
+    GradientBackend,
+    scipy_optimize_funcs_from_loss,
+)
+
+_log = logging.getLogger(__name__)
+
+
+def get_conditional_gaussian_approximation(
+    x: TensorVariable,
+    Q: TensorVariable | ArrayLike,
+    mu: TensorVariable | ArrayLike,
+    args: list[TensorVariable] | None = None,
+    model: pm.Model | None = None,
+    method: minimize_method = "BFGS",
+    use_jac: bool = True,
+    use_hess: bool = False,
+    optimizer_kwargs: dict | None = None,
+) -> Callable:
+    """
+    Returns a function to estimate the a posteriori log probability of a latent Gaussian field x and its mode x0 using the Laplace approximation.
+
+    That is:
+    y | x, sigma ~ N(Ax, sigma^2 W)
+    x | params ~ N(mu, Q(params)^-1)
+
+    We seek to estimate log(p(x | y, params)):
+
+    log(p(x | y, params)) = log(p(y | x, params)) + log(p(x | params)) + const
+
+    Let f(x) = log(p(y | x, params)). From the definition of our model above, we have log(p(x | params)) = -0.5*(x - mu).T Q (x - mu) + 0.5*logdet(Q).
+
+    This gives log(p(x | y, params)) = f(x) - 0.5*(x - mu).T Q (x - mu) + 0.5*logdet(Q). We will estimate this using the Laplace approximation by Taylor expanding f(x) about the mode.
+
+    Thus:
+
+    1. Maximize log(p(x | y, params)) = f(x) - 0.5*(x - mu).T Q (x - mu) wrt x (note that logdet(Q) does not depend on x) to find the mode x0.
+
+    2. Substitute x0 into the Laplace approximation expanded about the mode: log(p(x | y, params)) ~= -0.5*x.T (-f''(x0) + Q) x + x.T (Q.mu + f'(x0) - f''(x0).x0) + 0.5*logdet(Q).
+
+    Parameters
+    ----------
+    x: TensorVariable
+        The parameter with which to maximize wrt (that is, find the mode in x). In INLA, this is the latent field x~N(mu,Q^-1).
+    Q: TensorVariable | ArrayLike
+        The precision matrix of the latent field x.
+    mu: TensorVariable | ArrayLike
+        The mean of the latent field x.
+    args: list[TensorVariable]
+        Args to supply to the compiled function. That is, (x0, logp) = f(x, *args). If set to None, assumes the model RVs are args.
+    model: Model
+        PyMC model to use.
+    method: minimize_method
+        Which minimization algorithm to use.
+    use_jac: bool
+        If true, the minimizer will compute the gradient of log(p(x | y, params)).
+    use_hess: bool
+        If true, the minimizer will compute the Hessian log(p(x | y, params)).
+    optimizer_kwargs: dict
+        Kwargs to pass to scipy.optimize.minimize.
+
+    Returns
+    -------
+    f: Callable
+        A function which accepts a value of x and args and returns [x0, log(p(x | y, params))], where x0 is the mode. x is currently both the point at which to evaluate logp and the initial guess for the minimizer.
+    """
+    model = pm.modelcontext(model)
+
+    if args is None:
+        args = model.continuous_value_vars + model.discrete_value_vars
+
+    # f = log(p(y | x, params))
+    f_x = model.logp()
+    jac = pytensor.gradient.grad(f_x, x)
+    hess = pytensor.gradient.jacobian(jac.flatten(), x)
+
+    # log(p(x | y, params)) only including terms that depend on x for the minimization step (logdet(Q) ignored as it is a constant wrt x)
+    log_x_posterior = f_x - 0.5 * (x - mu).T @ Q @ (x - mu)
+
+    # Maximize log(p(x | y, params)) wrt x to find mode x0
+    x0, _ = minimize(
+        objective=-log_x_posterior,
+        x=x,
+        method=method,
+        jac=use_jac,
+        hess=use_hess,
+        optimizer_kwargs=optimizer_kwargs,
+    )
+
+    # require f'(x0) and f''(x0) for Laplace approx
+    jac = pytensor.graph.replace.graph_replace(jac, {x: x0})
+    hess = pytensor.graph.replace.graph_replace(hess, {x: x0})
+
+    # Full log(p(x | y, params)) using the Laplace approximation (up to a constant)
+    _, logdetQ = pt.nlinalg.slogdet(Q)
+    conditional_gaussian_approx = (
+        -0.5 * x.T @ (-hess + Q) @ x + x.T @ (Q @ mu + jac - hess @ x0) + 0.5 * logdetQ
+    )
+
+    # Currently x is passed both as the query point for f(x, args) = logp(x | y, params) AND as an initial guess for x0. This may cause issues if the query point is
+    # far from the mode x0 or in a neighbourhood which results in poor convergence.
+    return pytensor.function(args, [x0, conditional_gaussian_approx])
+
+
+def _unconstrained_vector_to_constrained_rvs(model):
+    outputs = get_default_varnames(model.unobserved_value_vars, include_transformed=True)
+    constrained_names = [
+        x.name for x in get_default_varnames(model.unobserved_value_vars, include_transformed=False)
+    ]
+    names = [x.name for x in outputs]
+
+    unconstrained_names = [name for name in names if name not in constrained_names]
+
+    new_outputs, unconstrained_vector = join_nonshared_inputs(
+        model.initial_point(),
+        inputs=model.value_vars,
+        outputs=outputs,
+    )
+
+    constrained_rvs = [x for x, name in zip(new_outputs, names) if name in constrained_names]
+    value_rvs = [x for x in new_outputs if x not in constrained_rvs]
+
+    unconstrained_vector.name = "unconstrained_vector"
+
+    # Redo the names list to ensure it is sorted to match the return order
+    names = [*constrained_names, *unconstrained_names]
+
+    return names, constrained_rvs, value_rvs, unconstrained_vector
+
+
+def model_to_laplace_approx(
+    model: pm.Model, unpacked_variable_names: list[str], chains: int = 1, draws: int = 500
+):
+    initial_point = model.initial_point()
+    raveled_vars = DictToArrayBijection.map(initial_point)
+    raveled_shape = raveled_vars.data.shape[0]
+
+    # temp_chain and temp_draw are a hack to allow sampling from the Laplace approximation. We only have one mu and cov,
+    # so we add batch dims (which correspond to chains and draws). But the names "chain" and "draw" are reserved.
+    names, constrained_rvs, value_rvs, unconstrained_vector = (
+        _unconstrained_vector_to_constrained_rvs(model)
+    )
+
+    coords = model.coords | {
+        "temp_chain": np.arange(chains),
+        "temp_draw": np.arange(draws),
+        "unpacked_variable_names": unpacked_variable_names,
+    }
+
+    with pm.Model(coords=coords, model=None) as laplace_model:
+        mu = pm.Flat("mean_vector", shape=(raveled_shape,))
+        cov = pm.Flat("covariance_matrix", shape=(raveled_shape, raveled_shape))
+        laplace_approximation = pm.MvNormal(
+            "laplace_approximation",
+            mu=mu,
+            cov=cov,
+            dims=["temp_chain", "temp_draw", "unpacked_variable_names"],
+            method="svd",
+        )
+
+        cast_to_var = partial(type_cast, Variable)
+        batched_rvs = vectorize_graph(
+            type_cast(list[Variable], constrained_rvs),
+            replace={cast_to_var(unconstrained_vector): cast_to_var(laplace_approximation)},
+        )
+
+        for name, batched_rv in zip(names, batched_rvs):
+            batch_dims = ("temp_chain", "temp_draw")
+            if batched_rv.ndim == 2:
+                dims = batch_dims
+            elif name in model.named_vars_to_dims:
+                dims = (*batch_dims, *model.named_vars_to_dims[name])
+            else:
+                dims = (*batch_dims, *[f"{name}_dim_{i}" for i in range(batched_rv.ndim - 2)])
+                initval = initial_point.get(name, None)
+                dim_shapes = initval.shape if initval is not None else batched_rv.type.shape[2:]
+                laplace_model.add_coords(
+                    {name: np.arange(shape) for name, shape in zip(dims[2:], dim_shapes)}
+                )
+
+            pm.Deterministic(name, batched_rv, dims=dims)
+
+    return laplace_model
+
+
+def unstack_laplace_draws(laplace_data, model, chains=2, draws=500):
+    """
+    The `model_to_laplace_approx` function returns a model with a single MvNormal distribution, draws from which are
+    in the unconstrained variable space. These might be interesting to the user, but since they come back stacked in a
+    single vector, it's not easy to work with.
+
+    This function unpacks each component of the vector into its own DataArray, with the appropriate dimensions and
+    coordinates, where possible.
+    """
+    initial_point = DictToArrayBijection.map(model.initial_point())
+
+    cursor = 0
+    unstacked_laplace_draws = {}
+    coords = model.coords | {"chain": range(chains), "draw": range(draws)}
+
+    # There are corner cases where the value_vars will not have the same dimensions as the random variable (e.g.
+    # simplex transform of a Dirichlet). In these cases, we don't try to guess what the labels should be, and just
+    # add an arviz-style default dim and label.
+    for rv, (name, shape, size, dtype) in zip(model.free_RVs, initial_point.point_map_info):
+        rv_dims = []
+        for i, dim in enumerate(
+            model.named_vars_to_dims.get(rv.name, [f"{name}_dim_{i}" for i in range(len(shape))])
+        ):
+            if coords.get(dim) and shape[i] == len(coords[dim]):
+                rv_dims.append(dim)
+            else:
+                rv_dims.append(f"{name}_dim_{i}")
+                coords[f"{name}_dim_{i}"] = np.arange(shape[i])
+
+        dims = ("chain", "draw", *rv_dims)
+
+        values = (
+            laplace_data[..., cursor : cursor + size].reshape((chains, draws, *shape)).astype(dtype)
+        )
+        unstacked_laplace_draws[name] = xr.DataArray(
+            values, dims=dims, coords={dim: list(coords[dim]) for dim in dims}
+        )
+
+        cursor += size
+
+    unstacked_laplace_draws = xr.Dataset(unstacked_laplace_draws)
+
+    return unstacked_laplace_draws
+
+
+def fit_laplace(
+    optimize_method: minimize_method | Literal["basinhopping"] = "BFGS",
+    *,
+    model: pm.Model | None = None,
+    use_grad: bool | None = None,
+    use_hessp: bool | None = None,
+    use_hess: bool | None = None,
+    initvals: dict | None = None,
+    random_seed: int | np.random.Generator | None = None,
+    jitter_rvs: list[pt.TensorVariable] | None = None,
+    progressbar: bool = True,
+    include_transformed: bool = True,
+    gradient_backend: GradientBackend = "pytensor",
+    chains: int = 2,
+    draws: int = 500,
+    optimizer_kwargs: dict | None = None,
+    compile_kwargs: dict | None = None,
+) -> az.InferenceData:
+    """
+    Create a Laplace (quadratic) approximation for a posterior distribution.
+
+    This function generates a Laplace approximation for a given posterior distribution using a specified
+    number of draws. This is useful for obtaining a parametric approximation to the posterior distribution
+    that can be used for further analysis.
+
+    Parameters
+    ----------
+    model : pm.Model
+        The PyMC model to be fit. If None, the current model context is used.
+    optimize_method : str
+        The optimization method to use. Valid choices are: Nelder-Mead, Powell, CG, BFGS, L-BFGS-B, TNC, SLSQP,
+        trust-constr, dogleg, trust-ncg, trust-exact, trust-krylov, and basinhopping.
+
+        See scipy.optimize.minimize documentation for details.
+    use_grad : bool | None, optional
+        Whether to use gradients in the optimization. Defaults to None, which determines this automatically based on
+        the ``method``.
+    use_hessp : bool | None, optional
+        Whether to use Hessian-vector products in the optimization. Defaults to None, which determines this automatically based on
+        the ``method``.
+    use_hess : bool | None, optional
+        Whether to use the Hessian matrix in the optimization. Defaults to None, which determines this automatically based on
+        the ``method``.
+    initvals : None | dict, optional
+        Initial values for the model parameters, as str:ndarray key-value pairs. Paritial initialization is permitted.
+         If None, the model's default initial values are used.
+    random_seed : None | int | np.random.Generator, optional
+        Seed for the random number generator or a numpy Generator for reproducibility
+    jitter_rvs : list of TensorVariables, optional
+        Variables whose initial values should be jittered. If None, all variables are jittered.
+    progressbar : bool, optional
+        Whether to display a progress bar during optimization. Defaults to True.
+    include_transformed: bool, default True
+        Whether to include transformed variables in the output. If True, transformed variables will be included in the
+        output InferenceData object. If False, only the original variables will be included.
+    gradient_backend: str, default "pytensor"
+        The backend to use for gradient computations. Must be one of "pytensor" or "jax".
+    chains: int, default: 2
+        The number of chain dimensions to sample. Note that this is *not* the number of chains to run in parallel,
+        because the Laplace approximation is not an MCMC method. This argument exists to ensure that outputs are
+        compatible with the ArviZ library.
+    draws: int, default: 500
+        The number of samples to draw from the approximated posterior. Totals samples will be chains * draws.
+    optimizer_kwargs
+        Additional keyword arguments to pass to the ``scipy.optimize`` function being used. Unless
+        ``method = "basinhopping"``, ``scipy.optimize.minimize`` will be used. For ``basinhopping``,
+        ``scipy.optimize.basinhopping`` will be used. See the documentation of these functions for details.
+    compile_kwargs: dict, optional
+        Additional keyword arguments to pass to pytensor.function.
+
+    Returns
+    -------
+    :class:`~arviz.InferenceData`
+        An InferenceData object containing the approximated posterior samples.
+
+    Examples
+    --------
+    >>> from pymc_extras.inference import fit_laplace
+    >>> import numpy as np
+    >>> import pymc as pm
+    >>> import arviz as az
+    >>> y = np.array([2642, 3503, 4358]*10)
+    >>> with pm.Model() as m:
+    >>>     logsigma = pm.Uniform("logsigma", 1, 100)
+    >>>     mu = pm.Uniform("mu", -10000, 10000)
+    >>>     yobs = pm.Normal("y", mu=mu, sigma=pm.math.exp(logsigma), observed=y)
+    >>>     idata = fit_laplace()
+
+    Notes
+    -----
+    This method of approximation may not be suitable for all types of posterior distributions,
+    especially those with significant skewness or multimodality.
+
+    See Also
+    --------
+    fit : Calling the inference function 'fit' like pmx.fit(method="laplace", model=m)
+          will forward the call to 'fit_laplace'.
+
+    """
+    compile_kwargs = {} if compile_kwargs is None else compile_kwargs
+    optimizer_kwargs = {} if optimizer_kwargs is None else optimizer_kwargs
+    model = pm.modelcontext(model) if model is None else model
+
+    idata = find_MAP(
+        method=optimize_method,
+        model=model,
+        use_grad=use_grad,
+        use_hessp=use_hessp,
+        use_hess=use_hess,
+        initvals=initvals,
+        random_seed=random_seed,
+        jitter_rvs=jitter_rvs,
+        progressbar=progressbar,
+        include_transformed=include_transformed,
+        gradient_backend=gradient_backend,
+        compile_kwargs=compile_kwargs,
+        **optimizer_kwargs,
+    )
+
+    unpacked_variable_names = idata.fit["mean_vector"].coords["rows"].values.tolist()
+
+    if "covariance_matrix" not in idata.fit:
+        # The user didn't use `use_hess` or `use_hessp` (or an optimization method that returns an inverse Hessian), so
+        # we have to go back and compute the Hessian at the MAP point now.
+        frozen_model = freeze_dims_and_data(model)
+        initial_params = _make_initial_point(frozen_model, initvals, random_seed, jitter_rvs)
+
+        _, f_hessp = scipy_optimize_funcs_from_loss(
+            loss=-frozen_model.logp(jacobian=False),
+            inputs=frozen_model.continuous_value_vars + frozen_model.discrete_value_vars,
+            initial_point_dict=DictToArrayBijection.rmap(initial_params),
+            use_grad=False,
+            use_hess=False,
+            use_hessp=True,
+            gradient_backend=gradient_backend,
+            compile_kwargs=compile_kwargs,
+        )
+        H_inv = _compute_inverse_hessian(
+            optimizer_result=None,
+            optimal_point=idata.fit.mean_vector.values,
+            f_fused=None,
+            f_hessp=f_hessp,
+            use_hess=False,
+            method=optimize_method,
+        )
+
+        idata.fit["covariance_matrix"] = xr.DataArray(
+            H_inv,
+            dims=("rows", "columns"),
+            coords={"rows": unpacked_variable_names, "columns": unpacked_variable_names},
+        )
+
+    with model_to_laplace_approx(model, unpacked_variable_names, chains, draws) as laplace_model:
+        new_posterior = (
+            pm.sample_posterior_predictive(
+                idata.fit.expand_dims(chain=[0], draw=[0]),
+                extend_inferencedata=False,
+                random_seed=random_seed,
+                var_names=[
+                    "laplace_approximation",
+                    *[x.name for x in laplace_model.deterministics],
+                ],
+            )
+            .posterior_predictive.squeeze(["chain", "draw"])
+            .drop_vars(["chain", "draw"])
+            .rename({"temp_chain": "chain", "temp_draw": "draw"})
+        )
+
+        if include_transformed:
+            idata.unconstrained_posterior = unstack_laplace_draws(
+                new_posterior.laplace_approximation.values, model, chains=chains, draws=draws
+            )
+
+        idata.posterior = new_posterior.drop_vars(
+            ["laplace_approximation", "unpacked_variable_names"]
+        )
+
+    return idata

pymc_extras/inference/laplace_approx/scipy_interface.py

@@ -0,0 +1,242 @@
+from collections.abc import Callable
+from importlib.util import find_spec
+from typing import Literal, get_args
+
+import numpy as np
+import pymc as pm
+import pytensor
+
+from pymc import join_nonshared_inputs
+from pytensor import tensor as pt
+from pytensor.compile import Function
+from pytensor.tensor import TensorVariable
+
+GradientBackend = Literal["pytensor", "jax"]
+VALID_BACKENDS = get_args(GradientBackend)
+
+
+def _compile_grad_and_hess_to_jax(
+    f_fused: Function, use_hess: bool, use_hessp: bool
+) -> tuple[Callable | None, Callable | None]:
+    """
+    Compile loss function gradients using JAX.
+
+    Parameters
+    ----------
+    f_fused: Function
+        The loss function to compile gradients for. Expected to be a pytensor function that returns a scalar loss,
+        compiled with mode="JAX".
+    use_hess: bool
+        Whether to compile a function to compute the hessian of the loss function.
+    use_hessp: bool
+        Whether to compile a function to compute the hessian-vector product of the loss function.
+
+    Returns
+    -------
+    f_fused: Callable
+        The compiled loss function and gradient function, which may also compute the hessian if requested.
+    f_hessp: Callable | None
+        The compiled hessian-vector product function, or None if use_hessp is False.
+    """
+    import jax
+
+    f_hessp = None
+
+    orig_loss_fn = f_fused.vm.jit_fn
+
+    if use_hess:
+
+        @jax.jit
+        def loss_fn_fused(x):
+            loss_and_grad = jax.value_and_grad(lambda x: orig_loss_fn(x)[0])(x)
+            hess = jax.hessian(lambda x: orig_loss_fn(x)[0])(x)
+            return *loss_and_grad, hess
+
+    else:
+
+        @jax.jit
+        def loss_fn_fused(x):
+            return jax.value_and_grad(lambda x: orig_loss_fn(x)[0])(x)
+
+    if use_hessp:
+
+        def f_hessp_jax(x, p):
+            y, u = jax.jvp(lambda x: loss_fn_fused(x)[1], (x,), (p,))
+            return jax.numpy.stack(u)
+
+        f_hessp = jax.jit(f_hessp_jax)
+
+    return loss_fn_fused, f_hessp
+
+
+def _compile_functions_for_scipy_optimize(
+    loss: TensorVariable,
+    inputs: list[TensorVariable],
+    compute_grad: bool,
+    compute_hess: bool,
+    compute_hessp: bool,
+    compile_kwargs: dict | None = None,
+) -> list[Function] | list[Function, Function | None, Function | None]:
+    """
+    Compile loss functions for use with scipy.optimize.minimize.
+
+    Parameters
+    ----------
+    loss: TensorVariable
+        The loss function to compile.
+    inputs: list[TensorVariable]
+        A single flat vector input variable, collecting all inputs to the loss function. Scipy optimize routines
+        expect the function signature to be f(x, *args), where x is a 1D array of parameters.
+    compute_grad: bool
+        Whether to compile a function that computes the gradients of the loss function.
+    compute_hess: bool
+        Whether to compile a function that computes the Hessian of the loss function.
+    compute_hessp: bool
+        Whether to compile a function that computes the Hessian-vector product of the loss function.
+    compile_kwargs: dict, optional
+        Additional keyword arguments to pass to the ``pm.compile`` function.
+
+    Returns
+    -------
+    f_fused: Function
+        The compiled loss function, which may also include gradients and hessian if requested.
+    f_hessp: Function | None
+        The compiled hessian-vector product function, or None if compute_hessp is False.
+    """
+    compile_kwargs = {} if compile_kwargs is None else compile_kwargs
+
+    loss = pm.pytensorf.rewrite_pregrad(loss)
+    f_hessp = None
+
+    # In the simplest case, we only compile the loss function. Return it as a list to keep the return type consistent
+    # with the case where we also compute gradients, hessians, or hessian-vector products.
+    if not (compute_grad or compute_hess or compute_hessp):
+        f_loss = pm.compile(inputs, loss, **compile_kwargs)
+        return [f_loss]
+
+    # Otherwise there are three cases. If the user only wants the loss function and gradients, we compile a single
+    # fused function and return it. If the user also wants the hessian, the fused function will return the loss,
+    # gradients and hessian. If the user wants gradients and hess_p, we return a fused function that returns the loss
+    # and gradients, and a separate function for the hessian-vector product.
+
+    if compute_hessp:
+        # Handle this first, since it can be compiled alone.
+        p = pt.tensor("p", shape=inputs[0].type.shape)
+        hessp = pytensor.gradient.hessian_vector_product(loss, inputs, p)
+        f_hessp = pm.compile([*inputs, p], hessp[0], **compile_kwargs)
+
+    outputs = [loss]
+
+    if compute_grad:
+        grads = pytensor.gradient.grad(loss, inputs)
+        grad = pt.concatenate([grad.ravel() for grad in grads])
+        outputs.append(grad)
+
+    if compute_hess:
+        hess = pytensor.gradient.jacobian(grad, inputs)[0]
+        outputs.append(hess)
+
+    f_fused = pm.compile(inputs, outputs, **compile_kwargs)
+
+    return [f_fused, f_hessp]
+
+
+def scipy_optimize_funcs_from_loss(
+    loss: TensorVariable,
+    inputs: list[TensorVariable],
+    initial_point_dict: dict[str, np.ndarray | float | int],
+    use_grad: bool,
+    use_hess: bool,
+    use_hessp: bool,
+    gradient_backend: GradientBackend = "pytensor",
+    compile_kwargs: dict | None = None,
+) -> tuple[Callable, ...]:
+    """
+    Compile loss functions for use with scipy.optimize.minimize.
+
+    Parameters
+    ----------
+    loss: TensorVariable
+        The loss function to compile.
+    inputs: list[TensorVariable]
+        The input variables to the loss function.
+    initial_point_dict: dict[str, np.ndarray | float | int]
+        Dictionary mapping variable names to initial values. Used to determine the shapes of the input variables.
+    use_grad: bool
+        Whether to compile a function that computes the gradients of the loss function.
+    use_hess: bool
+        Whether to compile a function that computes the Hessian of the loss function.
+    use_hessp: bool
+        Whether to compile a function that computes the Hessian-vector product of the loss function.
+    gradient_backend: str, default "pytensor"
+        Which backend to use to compute gradients. Must be one of "jax" or "pytensor"
+    compile_kwargs:
+        Additional keyword arguments to pass to the ``pm.compile`` function.
+
+    Returns
+    -------
+    f_fused: Callable
+        The compiled loss function, which may also include gradients and hessian if requested.
+    f_hessp: Callable | None
+        The compiled hessian-vector product function, or None if use_hessp is False.
+    """
+
+    compile_kwargs = {} if compile_kwargs is None else compile_kwargs
+
+    if use_hess and not use_grad:
+        raise ValueError("Cannot compute hessian without also computing the gradient")
+
+    if gradient_backend not in VALID_BACKENDS:
+        raise ValueError(
+            f"Invalid gradient backend: {gradient_backend}. Must be one of {VALID_BACKENDS}"
+        )
+
+    use_jax_gradients = (gradient_backend == "jax") and use_grad
+    if use_jax_gradients and not find_spec("jax"):
+        raise ImportError("JAX must be installed to use JAX gradients")
+
+    mode = compile_kwargs.get("mode", None)
+    if mode is None and use_jax_gradients:
+        compile_kwargs["mode"] = "JAX"
+    elif mode != "JAX" and use_jax_gradients:
+        raise ValueError(
+            'jax gradients can only be used when ``compile_kwargs["mode"]`` is set to "JAX"'
+        )
+
+    if not isinstance(inputs, list):
+        inputs = [inputs]
+
+    [loss], flat_input = join_nonshared_inputs(
+        point=initial_point_dict, outputs=[loss], inputs=inputs
+    )
+
+    # If we use pytensor gradients, we will use the pytensor function wrapper that handles shared variables. When
+    # computing jax gradients, we discard the function wrapper, so we can't handle shared variables --> rewrite them
+    # away.
+    if use_jax_gradients:
+        from pymc.sampling.jax import _replace_shared_variables
+
+        [loss] = _replace_shared_variables([loss])
+
+    compute_grad = use_grad and not use_jax_gradients
+    compute_hess = use_hess and not use_jax_gradients
+    compute_hessp = use_hessp and not use_jax_gradients
+
+    funcs = _compile_functions_for_scipy_optimize(
+        loss=loss,
+        inputs=[flat_input],
+        compute_grad=compute_grad,
+        compute_hess=compute_hess,
+        compute_hessp=compute_hessp,
+        compile_kwargs=compile_kwargs,
+    )
+
+    # Depending on the requested functions, f_fused will either be the loss function, the loss function with gradients,
+    # or the loss function with gradients and hessian.
+    f_fused = funcs.pop(0)
+    f_hessp = funcs.pop(0) if compute_hessp else None
+
+    if use_jax_gradients:
+        f_fused, f_hessp = _compile_grad_and_hess_to_jax(f_fused, use_hess, use_hessp)
+
+    return f_fused, f_hessp