folx 0.2__py3-none-any.whl

folx/hessian.py

@@ -0,0 +1,346 @@
+import functools
+import logging
+from typing import Callable, Sequence
+
+import jax
+import jax.flatten_util as jfu
+import jax.numpy as jnp
+import jax.tree_util as jtu
+import jaxlib.xla_extension
+import numpy as np
+from jax import core
+
+from .api import (
+    JAC_DIM,
+    Array,
+    Axes,
+    CustomTraceJacHessianJac,
+    ExtraArgs,
+    ForwardFn,
+    FunctionFlags,
+    FwdJacobian,
+    FwdLaplArgs,
+    FwdLaplArray,
+    MergeFn,
+    PyTree,
+)
+from .utils import (
+    add_vmap_jacobian_dim,
+    array_wise_flat_wrap,
+    flat_wrap,
+    get_reduced_jacobians,
+    jac_jacT,
+    trace_jac_jacT,
+    trace_of_product,
+    vmap_sequences_and_squeeze,
+)
+
+
+def general_jac_hessian_jac(fn: ForwardFn, args: FwdLaplArgs, materialize_idx: Array | None):
+    # It's conceptually easier to work with the flattened version of the
+    # Hessian, since we can then use einsum to compute the trace.
+    flat_fn = flat_wrap(fn, *args.x)
+    flat_x = jfu.ravel_pytree(args.x)[0]
+    # We have to decide on an order in which we execute tr(HJJ^T).
+    # H will be of shape NxDxD, J is DxK where N could potentially be D.
+    # We will do the following:
+    # if K >= D, we compute
+    # JJ^T first and then the trace.
+    # if D < K, we compute HJ first and then the trace.
+    # We should also flatten our gradient tensor to a 2D matrix where the first dimension
+    # is the x0 dim and the second dim is the input dim.
+    grads_2d = get_reduced_jacobians(*args.jacobian, idx=materialize_idx)
+    grad_2d = jnp.concatenate([x.T for x in grads_2d], axis=0)
+    D, K = grad_2d.shape
+    if K > D:
+        # jax.hessian uses Fwd on Reverse AD
+        flat_hessian = jax.hessian(flat_fn)(flat_x)
+        flat_out = trace_of_product(flat_hessian, grad_2d @ grad_2d.T)
+    elif D > K:
+        # Directly copmute the trace of tr(HJJ^T)=tr(J^THJ)
+        @functools.partial(jax.vmap, in_axes=-1, out_axes=-1)
+        def vhvp(tangent):
+            def vjp(x):
+                @functools.partial(jax.vmap, in_axes=(None, -1), out_axes=-1)
+                def jvp(x, tangent):
+                    return jax.jvp(flat_fn, (x,), (tangent,))[1]
+
+                return jvp(x, grad_2d)
+
+            return jax.jvp(vjp, (flat_x,), (tangent,))[1]
+
+        flat_out = jnp.trace(vhvp(grad_2d), axis1=-2, axis2=-1)
+    else:
+        # Implementation where we compute HJ and then the trace via
+        # the sum of hadamard product
+        @functools.partial(jax.vmap, in_axes=-1, out_axes=-1)
+        def hvp(tangent):
+            def jacobian(x):
+                return jax.jacrev(flat_fn)(x)
+
+            return jax.jvp(jacobian, (flat_x,), (tangent,))[1]
+
+        HJ = hvp(grad_2d)  # N x D x K
+        flat_out = trace_of_product(HJ, grad_2d)  # N x D x K and D x K
+
+    # since f(x) and nabla f(x) should have the same structure, we can use the
+    # structure of f(x) to unravel the flat_out
+    unravel = jfu.ravel_pytree(fn(*args.x))[1]
+    return unravel(flat_out)
+
+
+def off_diag_jac_hessian_jac(fn: ForwardFn, args: FwdLaplArgs, materialize_idx: Array | None):
+    # if we know that a function is linear in one arguments, it's hessian must be off diagonal
+    # thus we can safe some computation by only computing the off diagonal part of the hessian.
+    assert len(args) == 2, "Off diag hessian only supports 2 args at the moment."
+
+    def flat_arr(x: FwdLaplArray) -> Array:
+        return jfu.ravel_pytree(x.x)[0]
+
+    flat_fn = array_wise_flat_wrap(fn, *args.x)
+
+    def jac_lhs(lhs, rhs):
+        return jax.jacobian(flat_fn, argnums=0)(lhs, rhs)
+
+    hessian = jax.jacobian(jac_lhs, argnums=1)(flat_arr(args.arrays[0]), flat_arr(args.arrays[1]))
+
+    flat_out = 2 * trace_of_product(
+        hessian, jac_jacT(args.arrays[0].jacobian, args.arrays[1].jacobian, materialize_idx)
+    )
+    unravel = jfu.ravel_pytree(fn(*args.x))[1]
+    return unravel(flat_out)
+
+
+def mul_jac_hessian_jac(fn: ForwardFn, args: FwdLaplArgs, shared_idx: Array | None):
+    # For a dot product we know that the hessian looks like this:
+    # [0, I]
+    # [I, 0]
+    # where I is the identity matrix of the same shape as the input.
+    assert len(args) == 2, "Dot product only supports two args."
+    flat_out = (
+        2 * trace_jac_jacT(args.arrays[0].jacobian, args.arrays[1].jacobian, shared_idx)[None]
+    )
+    unravel = jfu.ravel_pytree(fn(*args.x))[1]
+    return unravel(flat_out)
+
+
+def remove_fill(arrs: np.ndarray, find_unique: bool = False):
+    """
+    Remove the fill value from an array. As the tensors might not be shaped correctly
+    afterwards, we reduce all the leading dimensions by lists.
+
+    Args:
+        - arrs: array to remove fill value from
+    Returns:
+        - arrs: nested lists of arrays without fill value
+    """
+    if arrs.size == 0:
+        return arrs
+    if arrs[0].ndim >= 1:
+        return [remove_fill(x, find_unique=find_unique) for x in arrs]
+    if find_unique:
+        arrs = np.unique(arrs)
+    return arrs[arrs >= 0]  # type: ignore
+
+
+def merge_and_populate(arrs: Sequence[np.ndarray], operation: Callable[[np.ndarray, np.ndarray], np.ndarray]):
+    """
+    The arrays are assumed to be of the same shape. We look at the intersection of all arrays.
+    We then find the maximum intersection size and fill all arrays to that size.
+
+    Args:
+        - arrs: list of arrays
+    Returns:
+        - arrs: np.ndarray where only intersections are kept and all arrays are filled to the same size.
+    """
+    result = jtu.tree_map(
+        lambda *x: functools.reduce(operation, tuple(x[1:]), x[0]),
+        *arrs,
+        is_leaf=lambda x: isinstance(x, np.ndarray)
+    )
+    sizes = jtu.tree_map(lambda x: x.size, result, is_leaf=lambda x: isinstance(x, np.ndarray))
+    max_size = np.max(jtu.tree_leaves(sizes))
+    result = jtu.tree_map(
+        lambda x: np.concatenate([x, np.full(max_size - x.size, -1, dtype=x.dtype)]),
+        result,
+        is_leaf=lambda x: isinstance(x, np.ndarray),
+    )
+    return np.asarray(result, dtype=int)
+
+
+def find_materialization_idx(lapl_args: FwdLaplArgs, in_axes, flags: FunctionFlags, threshold: int):
+    if not lapl_args.any_jacobian_weak:
+        return None
+    # TODO: Rewrite this!! This is quity messy and inefficient.
+    # it assumes that we're only interested in the last dimension.
+    with core.new_main(core.EvalTrace, dynamic=True):
+        vmap_seq, (inp,) = vmap_sequences_and_squeeze(
+            ([j.mask for j in lapl_args.jacobian],),
+            ([j for j in add_vmap_jacobian_dim(lapl_args, FwdLaplArgs(in_axes)).jacobian],),
+        )
+        max_size = np.max([np.sum(j.unique_idx >= 0, dtype=int) for j in lapl_args.jacobian])
+        # This can be quite memory intensive, so we try to do it on the GPU and
+        # if that fails we just use the CPU. On the CPU this takes quite some time.
+        # TODO: work on a more memory efficient implementation!
+        unique_fn = functools.partial(jnp.unique, size=max_size + 1, fill_value=-1)
+
+        def idx_fn(x):
+            return jtu.tree_map(unique_fn, x)
+        for s in vmap_seq[::-1]:
+            idx_fn = jax.vmap(idx_fn, in_axes=s)
+        try:
+            # This path is more memory intensive by using the GPU to find uniques but
+            # potentially fails if the arrays are too large.
+            # +1 because we need to accomodate the -1.
+            arrs = np.asarray(idx_fn(inp), dtype=int)
+        except jaxlib.xla_extension.XlaRuntimeError:
+            logging.info(
+                "Failed to find unique elements on GPU, falling back to CPU. This will be slow."
+            )
+            with jax.default_device(jax.devices("cpu")[0]):
+                arrs = np.asarray(idx_fn(inp), dtype=int)
+        filtered_arrs = remove_fill(arrs, False)
+
+    if FunctionFlags.LINEAR_IN_ONE in flags:
+        # For off diagonal Hessians we only need to look at the intersection between
+        # all arrays rather than their union.
+        idx = merge_and_populate(filtered_arrs, np.intersect1d) # type: ignore
+    else:
+        idx = merge_and_populate(filtered_arrs, np.union1d) # type: ignore
+    idx = np.moveaxis(idx, -1, JAC_DIM)
+
+    if idx.shape[JAC_DIM] >= max_size or idx.shape[JAC_DIM] > threshold:
+        idx = None
+    return idx
+
+
+def remove_zero_entries(lapl_args: FwdLaplArgs, materialize_idx: np.ndarray | None):
+    if materialize_idx is None:
+        return lapl_args, None, None
+    
+    mask = (materialize_idx != -1).any(0)
+    if mask.sum() > 0.5 * mask.size:
+        # this is a heuristic to avoid having unnecessary indexing overhead for
+        # insufficiently sparse masks.
+        return lapl_args, materialize_idx, None
+
+    indices = np.where(mask)
+    new_mat_idx = materialize_idx[(slice(None), *indices)]
+    new_arrs = []
+    for arg in lapl_args.arrays:
+        brdcast_dims = np.where(np.array(arg.x.shape) == 1)[0]
+        idx = tuple(
+            0 if i in brdcast_dims else x
+            for i, x in enumerate(indices)
+        )
+        new_arrs.append(FwdLaplArray(
+            x=arg.x[idx],
+            jacobian=FwdJacobian(
+                data=arg.jacobian.data[(slice(None), *idx)],
+                x0_idx=arg.jacobian.x0_idx[(slice(None), *idx)], # type: ignore
+            ),
+            laplacian=arg.laplacian[idx],
+        ))
+    new_args = FwdLaplArgs(tuple(new_arrs))
+    return new_args, new_mat_idx, mask
+
+
+def vmapped_jac_hessian_jac(
+    fwd: ForwardFn,
+    flags: FunctionFlags,
+    custom_jac_hessian_jac: CustomTraceJacHessianJac | None,
+    extra_args: ExtraArgs,
+    in_axes: Axes,
+    extra_in_axes: Axes,
+    merge: MergeFn,
+    sparsity_threshold: int,
+    lapl_args: FwdLaplArgs,
+) -> PyTree[Array]:
+    # Determine output structure
+    def merged_fn(*x: Array):
+        return fwd(*merge(x, extra_args))
+    out = merged_fn(*lapl_args.x)
+    unravel = jfu.ravel_pytree(out)[1]
+
+    materialize_idx = find_materialization_idx(lapl_args, in_axes, flags, sparsity_threshold)
+    if materialize_idx is None:
+        lapl_args = lapl_args.dense
+    if materialize_idx is not None and materialize_idx.shape[JAC_DIM] == 0:
+        return jnp.zeros(())
+    
+    # If we do a dot product (not a hadamard product) we can check for empty hessian entries
+    if FunctionFlags.DOT_PRODUCT in flags and all(len(a) == 1 for a in in_axes):
+        lapl_args, materialize_idx, mask = remove_zero_entries(lapl_args, materialize_idx)
+        in_axes = jtu.tree_map(lambda _: -1, in_axes)
+    else:
+        mask = None
+
+    # Broadcast and flatten all arguments
+    vmap_seq, (lapl_args, extra_args) = vmap_sequences_and_squeeze(
+        (lapl_args, extra_args),
+        (add_vmap_jacobian_dim(lapl_args, FwdLaplArgs(in_axes)), extra_in_axes),
+    )
+    # Hessian computation
+    def hess_transform(args: FwdLaplArgs, extra_args: ExtraArgs, materialize_idx):
+        def merged_fn(*x):
+            return fwd(*merge(x, extra_args))
+        
+        if custom_jac_hessian_jac is not None:
+            result = custom_jac_hessian_jac(args, extra_args, merge, materialize_idx)
+        elif FunctionFlags.MULTIPLICATION in flags:
+            result = mul_jac_hessian_jac(merged_fn, args, materialize_idx)
+        elif FunctionFlags.LINEAR_IN_ONE in flags:
+            result = off_diag_jac_hessian_jac(merged_fn, args, materialize_idx)
+        else:
+            result = general_jac_hessian_jac(merged_fn, args, materialize_idx)
+        return result
+
+    # TODO: this implementation also assumes that we only reduce the last dimension.
+    for axes in vmap_seq[::-1]:
+        hess_transform = jax.vmap(
+            hess_transform, in_axes=(*axes, (None if materialize_idx is None else 1))
+        )
+    # flatten to 1D and then unravel to the original structure
+    if mask is None:
+        flat_out = jfu.ravel_pytree(hess_transform(lapl_args, extra_args, materialize_idx))[0]
+    else:
+        flat_out = hess_transform(lapl_args, extra_args, materialize_idx)
+        result = jnp.zeros_like(out).at[mask].set(flat_out) # type: ignore
+        flat_out = jfu.ravel_pytree(result)[0]
+    return unravel(flat_out)
+
+
+def get_jacobian_hessian_jacobian_trace(
+    fwd: ForwardFn,
+    flags: FunctionFlags,
+    custom_jac_hessian_jac: CustomTraceJacHessianJac | None,
+    extra_args: ExtraArgs,
+    in_axes: Axes,
+    extra_in_axes: Axes,
+    merge: MergeFn,
+):
+    def hessian_transform(args: FwdLaplArgs, sparsity_threshold: int):
+        if FunctionFlags.LINEAR in flags:
+            return jnp.zeros(())
+        elif FunctionFlags.LINEAR_IN_ONE in flags and len(args.arrays) == 1:
+            return jnp.zeros(())
+        elif (
+            FunctionFlags.LINEAR_IN_FIRST in flags
+            and jtu.tree_leaves(merge(args.x, extra_args))[0] is args.x[0]
+            and len(args.arrays) == 1
+        ):
+            return jnp.zeros(())
+        else:
+            return vmapped_jac_hessian_jac(
+                fwd=fwd,
+                flags=flags,
+                custom_jac_hessian_jac=custom_jac_hessian_jac,
+                extra_args=extra_args,
+                in_axes=in_axes,
+                extra_in_axes=extra_in_axes,
+                merge=merge,
+                sparsity_threshold=sparsity_threshold,
+                lapl_args=args,
+            )
+
+    return hessian_transform

folx/interpreter.py

@@ -0,0 +1,221 @@
+import logging
+from collections import defaultdict
+from typing import Callable, ParamSpec, Sequence, TypeVar
+
+import jax
+import jax.numpy as jnp
+import jax.tree_util as jtu
+import numpy as np
+from jax import core
+from jax.util import safe_map
+
+from .api import Array, ArrayOrFwdLaplArray, FwdJacobian, FwdLaplArray, PyTree
+from .utils import extract_jacobian_mask, ravel
+from .wrapped_functions import get_laplacian
+
+R = TypeVar("R", bound=PyTree[Array])
+P = ParamSpec("P")
+
+
+class JaxExprEnvironment:
+    # A simple environment that keeps track of the variables
+    # and frees them once they are no longer needed.
+    env: dict[core.Var, ArrayOrFwdLaplArray]
+    reference_counter: dict[core.Var, int]
+
+    def __init__(self, jaxpr: core.Jaxpr, consts: Sequence[Array], *args: ArrayOrFwdLaplArray):
+        self.env = {}
+        self.reference_counter = defaultdict(int)
+        for v in jaxpr.invars + jaxpr.constvars:
+            if isinstance(v, core.Literal):
+                continue
+            self.reference_counter[v] += 1
+        eqn: core.JaxprEqn
+        for eqn in jaxpr.eqns:
+            for v in eqn.invars:
+                if isinstance(v, core.Literal):
+                    continue
+                self.reference_counter[v] += 1
+        for v in jaxpr.outvars:
+            if isinstance(v, core.Literal):
+                continue
+            self.reference_counter[v] = np.iinfo(np.int32).max
+        self.write_many(jaxpr.constvars, consts)
+        self.write_many(jaxpr.invars, args)
+
+    def read(self, var: core.Atom) -> ArrayOrFwdLaplArray:
+        if isinstance(var, core.Literal):
+            return var.val
+        self.reference_counter[var] -= 1
+        result = self.env[var]
+        if self.reference_counter[var] == 0:
+            del self.env[var]
+            del self.reference_counter[var]
+        return result
+
+    def write(self, var: core.Var, val: ArrayOrFwdLaplArray):
+        if self.reference_counter[var] > 0:
+            self.env[var] = val
+
+    def read_many(self, vars: Sequence[core.Atom]) -> list[ArrayOrFwdLaplArray]:
+        return safe_map(self.read, vars)
+
+    def write_many(self, vars: Sequence[core.Var], vals: Sequence[ArrayOrFwdLaplArray]):
+        return safe_map(self.write, vars, vals)
+
+
+def eval_jaxpr_with_forward_laplacian(jaxpr: core.Jaxpr, consts, *args, sparsity_threshold: int):
+    enable_sparsity = sparsity_threshold > 0
+    env = JaxExprEnvironment(jaxpr, consts, *args)
+
+    def eval_scan(eqn: core.JaxprEqn, invals):
+        n_carry, n_const = eqn.params["num_carry"], eqn.params["num_consts"]
+        in_const, in_carry, in_inp = invals[:n_const], invals[n_const:n_carry+n_const], invals[n_const+n_carry:]
+        carry_merge = extract_jacobian_mask(in_carry)
+        assert all(isinstance(x, Array) for x in in_inp), "Scan does not support scanning over input depenedent tensors.\nPlease unroll the loop."
+        def wrapped(carry, x):
+            result = eval_jaxpr_with_forward_laplacian(
+                eqn.params['jaxpr'].jaxpr,
+                (),
+                *in_const,
+                *carry_merge(carry),
+                *x,
+                sparsity_threshold=sparsity_threshold
+            )
+            return result[:n_carry], result[n_carry:]
+        first_carry, first_y = wrapped(in_carry, jtu.tree_map(lambda x: x[0], in_inp))
+        # Check whether jacobian sparsity matches
+        for a, b in zip(in_carry, first_carry):
+            if type(a) != type(b):
+                raise TypeError(f"Type mismatch in scan: {type(a)} != {type(b)}")
+            if isinstance(a, FwdLaplArray):
+                if not np.all(a.jacobian.x0_idx == b.jacobian.x0_idx): # type: ignore
+                    raise ValueError("Jacobian sparsity mismatch in scan.")
+        carry, y = jax.lax.scan(
+            wrapped, # type: ignore
+            in_carry,
+            in_inp,
+            length=eqn.params['length'],
+            reverse=eqn.params['reverse'],
+            unroll=eqn.params['unroll'],
+        )
+        carry = [
+            a._replace(jacobian=a.jacobian._replace(x0_idx=b.jacobian.x0_idx)) # type: ignore
+            if isinstance(a, FwdLaplArray) else a
+            for a, b in zip(carry, first_carry)
+        ]
+        y = [
+            a._replace(jacobian=a.jacobian._replace(x0_idx=b.jacobian.x0_idx)) # type: ignore
+            if isinstance(a, FwdLaplArray) else a
+            for a, b in zip(y, first_y)
+        ]
+        return *carry, *y
+
+    def eval_pjit(eqn: core.JaxprEqn, invals):
+        name = eqn.params["name"]
+        if fn := get_laplacian(name):
+            # TODO: this is a bit incomplete, e.g., kwargs?
+            outvals = fn(*invals, sparsity_threshold=sparsity_threshold)
+            if isinstance(outvals, (FwdLaplArray, Array)):
+                outvals = [outvals]  # TODO: Figure out how to properly handle outvals
+            return outvals
+        sub_expr: core.ClosedJaxpr = eqn.params["jaxpr"]
+        return eval_jaxpr_with_forward_laplacian(
+            sub_expr.jaxpr, sub_expr.literals, *invals, sparsity_threshold=sparsity_threshold
+        )
+
+    def eval_laplacian(eqn: core.JaxprEqn, invals):
+        fn = get_laplacian(eqn.primitive, True)
+        return fn(*invals, **eqn.params, sparsity_threshold=sparsity_threshold)
+
+    for eqn in jaxpr.eqns:
+        invals = env.read_many(eqn.invars)
+        # Eval expression
+        if all(not isinstance(x, FwdLaplArray) for x in invals):
+            subfuns, bind_params = eqn.primitive.get_bind_params(eqn.params)
+            # If non of the inputs were dependent on an FwdLaplArray,
+            # we can just use the regular primitive. This will avoid
+            # omnistaging. While this could cost us some memory and speed,
+            # it gives us access to more variables during tracing.
+            # https://github.com/google/jax/pull/3370
+            if all(not isinstance(x, core.Tracer) for x in invals) and enable_sparsity:
+                try:
+                    with core.new_main(core.EvalTrace, dynamic=True):
+                        outvals = eqn.primitive.bind(*subfuns, *invals, **bind_params)
+                except Exception as e:
+                    logging.warning(
+                        f"Could not perform operation {eqn.primitive.name} in eager execution despite it only depending on non-input dependent values. "
+                        "We switch to tracing rather than eager execution. This may impact sparsity propagation.\n"
+                        f"{e}"
+                    )
+                    outvals = eqn.primitive.bind(*subfuns, *invals, **bind_params)
+            else:
+                outvals = eqn.primitive.bind(*subfuns, *invals, **bind_params)
+        elif eqn.primitive.name == "scan":
+            outvals = eval_scan(eqn, invals)
+        elif eqn.primitive.name == "pjit":
+            outvals = eval_pjit(eqn, invals)
+        else:
+            outvals = eval_laplacian(eqn, invals)
+
+        # unify output
+        if not eqn.primitive.multiple_results:
+            outvals = [outvals]
+        # save output
+        env.write_many(eqn.outvars, outvals)  # type: ignore
+    return env.read_many(jaxpr.outvars)
+
+
+
+def init_forward_laplacian_state(*x: PyTree[Array], sparsity: bool) -> PyTree[FwdLaplArray]:
+    """
+    Initialize forward Laplacian state from a PyTree of arrays.
+    """
+    x_flat, unravel = ravel(x)
+    jac = jtu.tree_map(jnp.ones_like, x)
+    jac_idx = unravel(np.arange(x_flat.shape[0]))
+    if sparsity:
+        jac = jtu.tree_map(lambda j, i: FwdJacobian(j[None], np.array(i)[None]), jac, jac_idx)
+    else:
+        jac = jax.vmap(unravel)(jnp.eye(len(x_flat)))
+        jac = jtu.tree_map(FwdJacobian.from_dense, jac)
+    lapl_x = jtu.tree_map(jnp.zeros_like, x)
+    return jtu.tree_map(FwdLaplArray, x, jac, lapl_x)
+
+
+def forward_laplacian(
+    fn: Callable[P, PyTree[Array]],
+    sparsity_threshold: int | float = 0,
+    disable_jit: bool = False
+) -> Callable[P, PyTree[FwdLaplArray]]:
+    """
+    This function takes a function and returns a function that computes the Laplacian of the function.
+    The returned function will be jitted by default as running it in eager execution will typically be a lot slower.
+
+    Args:
+        - fn: function to compute the Laplacian of
+        - sparsity_threshold: threshold for sparsity propagation.
+            If the number of non-zero elements in the input is larger than this threshold,we will not propagate sparsity.
+            If the value is between 0 and 1, it will be interpreted as a fraction of the total number of elements.
+            If the value is larger than 1, it will be interpreted as an absolute number of elements.
+            If enabling sparsity, we recommend relatively large values like 0.6 as frequent materializations are slow.
+    """
+    def wrapped(*args: P.args, **kwargs: P.kwargs):
+        closed_jaxpr = jax.make_jaxpr(fn)(*args, **kwargs)
+        flat_args = jtu.tree_leaves(args)
+        if 0 < sparsity_threshold < 1:
+            threshold = int(sparsity_threshold * sum(x.size for x in flat_args))
+        else:
+            threshold = int(sparsity_threshold)
+        lapl_args = init_forward_laplacian_state(*flat_args, sparsity=threshold > 0)
+        out = eval_jaxpr_with_forward_laplacian(
+            closed_jaxpr.jaxpr, closed_jaxpr.literals, *lapl_args, sparsity_threshold=threshold
+        )
+        if len(out) == 1:
+            return out[0]
+        return out
+    
+    if disable_jit:
+        return wrapped
+
+    return jax.jit(wrapped) # type: ignore