Solve profiling

Opt-in debugging log of the actions the sparse solvers take. See the solve profile guide for a walkthrough.

create_solve_profile

create_solve_profile() -> SolveProfile

Create a new, empty SolveProfile, for debugging the sparse solvers.

Equivalent to SolveProfile(). Enter it as a context manager to record the analyze, factor, refactor, solve, and iterative-refinement operations run inside the block:

profile = splineax.create_solve_profile()
with profile:
    solution, state = splineax.linear_solve(operator, vector, solver)
    state.release()
print(profile)

---

profile_solves

profile_solves(
    fn: Callable[_P, _ReturnT], *, enabled: bool = True
) -> Callable[_P, tuple[_ReturnT, SolveProfile | None]]

profile_solves(
    fn: None = None, *, enabled: bool = True
) -> Callable[
    [Callable[_P, _ReturnT]],
    Callable[_P, tuple[_ReturnT, SolveProfile | None]],
]

profile_solves(
    fn: Callable[_P, _ReturnT] | None = None,
    *,
    enabled: bool = True,
) -> (
    Callable[_P, tuple[_ReturnT, SolveProfile | None]]
    | Callable[
        [Callable[_P, _ReturnT]],
        Callable[_P, tuple[_ReturnT, SolveProfile | None]],
    ]
)

Wrap fn so every call profiles its solves into a fresh SolveProfile.

The wrapped function returns (result, profile) instead of fn's own return value. Use it bare as @profile_solves, or with arguments as @profile_solves(enabled=False). With enabled=False, every call runs fn directly and returns (result, None). This lets you switch profiling off in one place without changing the call sites.

Apply this directly to a jax.jit/equinox.filter_jit-decorated function, and always call it through the decorated name. Used this way, the very first call for a given input shape happens through this wrapper with a profile active, so that shape's compiled executable is guaranteed to carry its profiling hooks. The hooks are only built in at compile time. Calling the undecorated jit function separately, or through a second wrapper with enabled=False, will forfeit profiling for any shape it compiles first.

@splineax.profile_solves
@equinox.filter_jit
def solve(operator, vector, solver):
    solution, state = splineax.linear_solve(operator, vector, solver)
    state.release()
    return solution

solution, profile = solve(operator, vector, solver)

---

SolveProfile

SolveProfile()

An ordered log of the operations taken across one or more state-sequences.

Create one with create_solve_profile. SolveProfile is itself a context manager: entering it (with profile:) makes it the active profile on this thread, so every operation recorded while it is entered lands in it, and nothing recorded outside any with block, or while a different profile is entered, does. A profile can be entered more than once, accumulating further records each time.

records is the log, and sequences slices it into one list per state-sequence (each generic init/init_symbolic starts a new one). Printing a SolveProfile renders the indented tree, with factorizations built anew and factorizations reused shown in different colors.

sequences property

sequences: list[list[ProfileRecord]]

The records grouped into state-sequences, split at each generic init.

__enter__

__enter__() -> SolveProfile

Make this profile the active one on this thread, for every operation inside.

render

render(color: bool | None = None) -> str

Render the tree. color forces ANSI on/off, and None auto-detects a TTY.

---

ProfileRecord dataclass

ProfileRecord(
    operation: str,
    solver: str | None = None,
    inputs: Mapping[str, Any] = dict(),
    outputs: Mapping[str, Any] = dict(),
    order: tuple[int, ...] = (),
)

One recorded operation in a SolveProfile.

operation is the operation name. A generic operation (init, init_symbolic, update, compute, track, release) has solver None. A solver-specific operation has solver set to the solver type (KLU, Pardiso, Spsolve, IterativeRefinement) and operation set to the library function it ran (analyze, factor, solve_with_numeric, ...). inputs and outputs hold the fields rendered as operation[inputs] => (outputs). order is the sort key for program order, taken when the operation is traced.