API reference¶

class PardisoSolver:
    """Reuses a single Pardiso factorization across many solves.

    The sparsity pattern (indptr and indices) is fixed for the solver's
    lifetime. The three Pardiso stages are kept as separate calls so callers
    control exactly what work happens on each one:

    - analyze() runs the symbolic phase for the pattern. Calling it again on
      the same solver re-analyzes in place, freeing the numeric factorization
      and reusing the same native handle rather than allocating a second one.
    - factorize() runs the first numeric factorization, and requires a prior
      analyze().
    - refactorize() updates the numeric factorization for new values on the
      same pattern, and requires a prior factorize(). It runs the same
      Pardiso phase as factorize(); the separate name and precondition make
      the reuse explicit at the call site.
    - solve() solves against whatever factorization is currently stored, and
      requires a prior factorize().
    - refactor_and_solve() factorizes for new values and solves in one call,
      reusing the analysis and requiring only a prior analyze(). It keeps no
      reference to the values, so unlike factorize() plus solve() it is safe
      to call from inside a jitted function where the values and right-hand
      side are tracers.

    Pardiso's parameters are recomputed fresh on every native call rather
    than persisted, so an options overlay (see
    pardiso_mkl_jax.iparm.PardisoOption) passed to a single method applies
    only to that call. The constructor's own options argument is the way to
    set an overlay for the solver's whole lifetime: it is applied to every
    call, and a per-call options argument layers on top of it, winning on any
    entry both set.

    Each call also records its diagnostics (see PardisoDiagnostics),
    readable afterward from last_diagnostics. Every method also takes
    return_diagnostics, which hands them back directly instead. That is the
    form to use under jit, where last_diagnostics is unavailable; see its
    docstring.

    PardisoSolver may be used as a context manager, which releases its native
    memory on exit, but this is optional. The cache behind every handle is
    bounded (set by PARDISO_MKL_JAX_FACTOR_CACHE), and any factorization that is
    evicted or released is rebuilt on next use from the matrix the call carries.
    So a solver that is never closed leaks at most one cache slot rather than
    unbounded memory, and reusing it after close() still works, it just rebuilds
    once. Close it, or use the with-block, to free that slot promptly.

        with PardisoSolver(indptr, indices, matrix_type=MatrixType.REAL_NONSYMMETRIC) as solver:
            solver.analyze(values)
            solver.factorize(values)
            x = solver.solve(b)

    The same calls work without the with-block, and close() stays available to
    release early.
    """

    def __init__(self, indptr, indices, *, matrix_type: MatrixType, options: OptionsLike = None):
        check_matrix_type_supported(matrix_type)
        if indptr.dtype.name != "int32" or indices.dtype.name != "int32":
            raise TypeError("indptr and indices must have dtype int32.")
        self._indptr = indptr
        self._indices = check_upper_triangular(indptr, indices, matrix_type)
        self._matrix_type = matrix_type
        self._dimension = matrix_dimension(indptr)
        # Validated once here rather than on every call that merges it in.
        self._options = canonicalize_overlay(options)
        # The token is only obtained from analyze(), which allocates the native
        # factorization it names, so there is nothing to hold until then.
        self._handle: primitive.FactorizationToken | None = None
        self._values = None
        self._closed = False
        self._analyzed = False
        self._factorized = False
        self._last_diagnostics: PardisoDiagnostics | None = None
        # Effective SCALING and WEIGHTED_MATCHING at the last analyze, which
        # every later phase is checked against. See _check_pivot_settings.
        self._analysis_pivot_settings: dict[int, int] = {}

    def __enter__(self) -> PardisoSolver:
        return self

    def __exit__(self, exc_type, exc_value, traceback) -> None:
        self.close()

    def close(self) -> None:
        """Release the native factorization. Called automatically by __exit__.

        Safe to call at any time. A later use of the solver rebuilds what was
        released. It only reclaims memory without a wasted rebuild when it runs
        after the last solve, which is the case here since close runs eagerly in
        Python. See the advanced-usage guide on releasing explicitly.
        """
        if not self._closed:
            if self._handle is not None:
                primitive.release(self._handle)
            self._closed = True

    @property
    def last_diagnostics(self) -> PardisoDiagnostics | None:
        """Diagnostics from the most recent analyze, factorize, refactorize, or solve call.

        None before any call has been made. Only ever reflects a successful
        call: a Pardiso error raises instead of returning, so a failed call
        leaves this at whatever it was before that call.

        Also None after a call that ran under jit, since the diagnostics only
        exist as tracers there and storing one would leak it out of its trace.
        Pass return_diagnostics to read diagnostics from a traced call, which
        is what refactor_and_solve() does for every call, traced or not: it
        never updates this at all.
        """
        return self._last_diagnostics

    def _merge(self, options: OptionsLike) -> tuple[tuple[int, int], ...]:
        """Layer a per-call overlay on top of the solver-wide one, per-call winning."""
        return merge_overlays(self._options, options)

    def _pivot_settings(self, overlay: tuple[tuple[int, int], ...]) -> dict[int, int]:
        """The values SCALING and WEIGHTED_MATCHING will actually take for a call.

        An entry the overlay sets takes that value; anything else falls back
        to this package's default for the matrix type.
        """
        defaults = primitive.default_iparm(self._matrix_type)
        entries = dict(overlay)
        return {
            index: int(entries.get(index, defaults[index]))
            for index in (PardisoOption.SCALING, PardisoOption.WEIGHTED_MATCHING)
        }

    def _check_pivot_settings(self, overlay: tuple[tuple[int, int], ...], stage: str) -> None:
        """Reject a call whose scaling or matching disagrees with the analysis.

        Pardiso computes its scaling and matching during analysis and expects
        the same settings at every later phase. Nothing in the native layer
        enforces that: each handler rebuilds iparm from scratch, so a
        disagreement is silently accepted and produces a wrong answer rather
        than an error. Catching it here is the only place it shows up.
        """
        for index, value in self._pivot_settings(overlay).items():
            analysis_value = self._analysis_pivot_settings[index]
            if value != analysis_value:
                name = PardisoOption(index).name
                raise ValueError(
                    f"{stage} would run with {name} (iparm[{index}]) = {value}, but the "
                    f"analysis for this solver ran with {analysis_value}. Pardiso computes "
                    "scaling and matching during analysis and expects them unchanged "
                    "afterwards. Set this option on the PardisoSolver constructor so it "
                    "applies to every call, or re-run analyze() with the new value."
                )

    def _record_diagnostics(self, final_iparm) -> PardisoDiagnostics:
        """Decode diagnostics, store them on the solver, and return them.

        Stores None instead when final_iparm is a tracer, which it is whenever
        the call is running under jit. Keeping the decoded tracer would leak
        it out of its trace, so reading last_diagnostics afterwards would
        raise rather than return anything useful. The returned value is the
        real one either way, so return_diagnostics still works under jit.
        """
        diagnostics = PardisoDiagnostics.from_iparm(final_iparm)
        self._last_diagnostics = None if isinstance(final_iparm, jax.core.Tracer) else diagnostics
        return diagnostics

    def _check_usable(self) -> None:
        # The context manager is optional now that a released factorization
        # rebuilds itself, so only a genuine close() blocks further use.
        if self._closed:
            raise RuntimeError("PardisoSolver is closed and can no longer be used.")

    # The overloads on analyze, factorize, and refactorize are here so that
    # `diagnostics = solver.factorize(values, return_diagnostics=True)` types
    # as a plain PardisoDiagnostics for callers, rather than something
    # optional they have to narrow before reading a field off it.
    @overload
    def analyze(
        self,
        values,
        *,
        options: OptionsLike = None,
        return_diagnostics: Literal[False] = False,
    ) -> None: ...

    @overload
    def analyze(
        self, values, *, options: OptionsLike = None, return_diagnostics: Literal[True]
    ) -> PardisoDiagnostics: ...

    def analyze(
        self, values, *, options: OptionsLike = None, return_diagnostics: bool = False
    ) -> PardisoDiagnostics | None:
        """Run the symbolic analysis (fill-reducing ordering) for the stored pattern.

        Takes a representative values array because Pardiso's default
        heuristics for non-symmetric matrices, scaling and matching, look at
        the numeric values during analysis. The permutation and scaling this
        produces stay valid for a later factorize() call with different
        values on the same pattern, so this only needs to run once.

        Calling it again on the same solver re-analyzes in place: the existing
        numeric factorization is freed and the same native handle is reused,
        so no second handle is allocated and nothing extra needs releasing.
        factorize() must run again afterwards before any solve(), since the
        factorization the re-analysis discarded is the one solve() would have
        used.

        Must be called outside jit. It stores the native handle on the solver,
        and under jit that handle is a tracer, which would escape its trace.
        Callers who need the whole lifecycle inside a jitted function use the
        pardiso_mkl_jax.primitive functions and thread the handle themselves.

        Returns the call's PardisoDiagnostics if return_diagnostics is set,
        and None otherwise.
        """
        self._check_usable()
        check_csr_arrays(self._indptr, self._indices, values)
        overlay = self._merge(options)
        if self._handle is None:
            self._handle, final_iparm = primitive.analyze(
                self._indptr,
                self._indices,
                values,
                matrix_type=self._matrix_type,
                options=overlay,
            )
        else:
            # Cleared before the call, not after. Re-analysis frees the
            # existing factorization first thing, so if it then fails there is
            # no analysis and no factors left to use, and the solver has to
            # say so rather than report the state it had going in.
            self._analyzed = False
            self._factorized = False
            self._values = None
            self._handle, final_iparm = primitive.reanalyze(
                self._handle,
                self._indptr,
                self._indices,
                values,
                matrix_type=self._matrix_type,
                options=overlay,
            )
        self._analysis_pivot_settings = self._pivot_settings(overlay)
        self._analyzed = True
        diagnostics = self._record_diagnostics(final_iparm)
        return diagnostics if return_diagnostics else None

    @overload
    def factorize(
        self,
        values,
        *,
        options: OptionsLike = None,
        return_diagnostics: Literal[False] = False,
    ) -> None: ...

    @overload
    def factorize(
        self, values, *, options: OptionsLike = None, return_diagnostics: Literal[True]
    ) -> PardisoDiagnostics: ...

    def factorize(
        self, values, *, options: OptionsLike = None, return_diagnostics: bool = False
    ) -> PardisoDiagnostics | None:
        """Run the first numeric factorization for values. Requires a prior analyze().

        Returns the call's PardisoDiagnostics if return_diagnostics is set,
        and None otherwise. That is where perturbed_pivot_count lives, which
        is Pardiso's own report that it could not pivot cleanly and the
        factorization may be unusable.
        """
        self._check_usable()
        if not self._analyzed:
            raise RuntimeError("factorize() requires analyze() to have been called first.")
        diagnostics = self._run_numeric_factorization(values, options=options, stage="factorize()")
        self._factorized = True
        return diagnostics if return_diagnostics else None

    @overload
    def refactorize(
        self,
        values,
        *,
        options: OptionsLike = None,
        return_diagnostics: Literal[False] = False,
    ) -> None: ...

    @overload
    def refactorize(
        self, values, *, options: OptionsLike = None, return_diagnostics: Literal[True]
    ) -> PardisoDiagnostics: ...

    def refactorize(
        self, values, *, options: OptionsLike = None, return_diagnostics: bool = False
    ) -> PardisoDiagnostics | None:
        """Update the numeric factorization with new values on the same pattern.

        Requires a prior factorize(). Skips the analysis phase, which is the
        cheap path when only the matrix values change between solves.

        Returns the call's PardisoDiagnostics if return_diagnostics is set,
        and None otherwise.
        """
        self._check_usable()
        if not self._factorized:
            raise RuntimeError("refactorize() requires factorize() to have been called first.")
        diagnostics = self._run_numeric_factorization(
            values, options=options, stage="refactorize()"
        )
        return diagnostics if return_diagnostics else None

    def _run_numeric_factorization(
        self, values, *, options: OptionsLike, stage: str
    ) -> PardisoDiagnostics:
        check_csr_arrays(self._indptr, self._indices, values)
        overlay = self._merge(options)
        self._check_pivot_settings(overlay, stage)
        self._handle, final_iparm = primitive.factor(
            self._handle,
            self._indptr,
            self._indices,
            values,
            matrix_type=self._matrix_type,
            options=overlay,
        )
        self._values = values
        return self._record_diagnostics(final_iparm)

    def solve(
        self,
        right_hand_side,
        *,
        transpose: bool = False,
        options: OptionsLike = None,
        return_diagnostics: bool = False,
    ):
        """Solve against the current factorization. Requires a prior factorize().

        Solves against A^T instead of A when transpose is set, reusing the
        same factorization: no extra factorize() call is needed to switch
        between the two, and consecutive calls with different transpose
        values are safe.

        Returns just the solution by default, or (solution, PardisoDiagnostics)
        if return_diagnostics is set. The latter is the form to use under jit,
        where last_diagnostics stays None.
        """
        self._check_usable()
        if not self._factorized:
            raise RuntimeError("solve() requires factorize() to have been called first.")
        overlay = self._merge(options)
        self._check_pivot_settings(overlay, "solve()")
        stacked_right_hand_side = right_hand_side[None, :]
        solution, final_iparm = primitive.solve_stateful(
            self._handle,
            self._indptr,
            self._indices,
            self._values,
            stacked_right_hand_side,
            matrix_type=self._matrix_type,
            transpose=transpose,
            options=overlay,
        )
        diagnostics = self._record_diagnostics(final_iparm)
        if return_diagnostics:
            return solution[0], diagnostics
        return solution[0]

    def refactor_and_solve(
        self,
        values,
        right_hand_side,
        *,
        transpose: bool = False,
        options: OptionsLike = None,
        return_diagnostics: bool = False,
    ):
        """Factorize for values and solve in one call, reusing the analysis.

        Requires a prior analyze(). Runs the numeric factorization and the
        solve as one combined Pardiso step (phase 23), reusing the symbolic
        analysis rather than re-running it. Because it is a single FFI call it
        stays correct inside a jitted function, where a separate factorize()
        and solve() would not be ordered, and it keeps no reference to values
        on the solver, so values and right_hand_side may be tracers.

        Solves against A^T instead of A when transpose is set.

        This is the one method that does not record its diagnostics on
        last_diagnostics even when it runs eagerly, since keeping nothing at
        all on the solver is the whole point of it. Pass return_diagnostics to
        get them back as a second return value instead.
        """
        self._check_usable()
        if not self._analyzed:
            raise RuntimeError("refactor_and_solve() requires analyze() to have been called first.")
        check_csr_arrays(self._indptr, self._indices, values)
        overlay = self._merge(options)
        self._check_pivot_settings(overlay, "refactor_and_solve()")
        stacked_right_hand_side = right_hand_side[None, :]
        solution, final_iparm = primitive.factor_and_solve_stateful(
            self._handle,
            self._indptr,
            self._indices,
            values,
            stacked_right_hand_side,
            matrix_type=self._matrix_type,
            transpose=transpose,
            options=overlay,
        )
        if return_diagnostics:
            return solution[0], PardisoDiagnostics.from_iparm(final_iparm)
        return solution[0]

class MatrixType(enum.IntEnum):
    """Pardiso matrix type codes, matching the `mtype` parameter.

    Only the real-valued members can be used in this version of the package,
    since it works with float64 values throughout. The complex members are
    included so the full set of matrix types Pardiso itself supports is
    documented here, and raise NotImplementedError if selected.

    The members whose values are mathematically symmetric or Hermitian
    (REAL_SYMMETRIC_POSITIVE_DEFINITE, REAL_SYMMETRIC_INDEFINITE,
    COMPLEX_HERMITIAN_POSITIVE_DEFINITE, COMPLEX_HERMITIAN_INDEFINITE, and
    COMPLEX_SYMMETRIC) require the CSR arrays to hold only the upper
    triangle, including the diagonal, not the full matrix: passing the full
    matrix corrupts Pardiso's factorization instead of raising a clear
    error. check_upper_triangular enforces this.

    The structurally symmetric members (REAL_STRUCTURALLY_SYMMETRIC and
    COMPLEX_STRUCTURALLY_SYMMETRIC) only assume a symmetric sparsity
    pattern, not symmetric values, so they need the full matrix like the
    nonsymmetric members do.
    """

    REAL_STRUCTURALLY_SYMMETRIC = 1
    """Real values, symmetric sparsity pattern, no assumption on the values themselves."""

    REAL_SYMMETRIC_POSITIVE_DEFINITE = 2
    """Real, symmetric, and positive definite. The cheapest and most stable case to factor."""

    REAL_SYMMETRIC_INDEFINITE = -2
    """Real and symmetric, but not guaranteed positive definite."""

    COMPLEX_STRUCTURALLY_SYMMETRIC = 3
    """Complex values, symmetric sparsity pattern. Not yet supported: requires complex values."""

    COMPLEX_HERMITIAN_POSITIVE_DEFINITE = 4
    """Complex and Hermitian positive definite. Not yet supported: requires complex values."""

    COMPLEX_HERMITIAN_INDEFINITE = -4
    """Complex and Hermitian, not guaranteed positive definite. Not yet supported."""

    COMPLEX_SYMMETRIC = 6
    """Complex and symmetric. Not yet supported: requires complex values."""

    REAL_NONSYMMETRIC = 11
    """Real, with no assumed symmetry. The general-purpose default matrix type."""

    COMPLEX_NONSYMMETRIC = 13
    """Complex, with no assumed symmetry. Not yet supported: requires complex values."""

class PardisoOption(enum.IntEnum):
    """Settable Pardiso iparm entries, one member per documented input option.

    Values are the 0-based iparm index. Reserved entries (which must stay 0)
    and pure-output entries (which Pardiso only ever writes, never reads) are
    intentionally not members here: reserved entries have no reason to be
    touched, and outputs are read back through PardisoDiagnostics instead.

    A few members need special handling in canonicalize_overlay beyond what
    their docstring alone conveys:

    - USE_DEFAULT_VALUES and WEIGHTED_MATCHING can be overridden freely, with
      no runtime warning: setting USE_DEFAULT_VALUES away from 1 is what a
      confirmed MKL segfault in weighted matching was worked around by
      forcing to 1 in the first place (see the "Solver settings" docs), and
      WEIGHTED_MATCHING is the actual crash-triggering setting that
      workaround exists to keep disabled, so a caller touching either one is
      knowingly reaching for the same danger.
    - INDEXING_STYLE can be overridden, but raises a runtime warning: every
      CSR array this package builds is zero-based, so anything else risks
      Pardiso silently misreading indptr and indices.
    - USER_PERMUTATION and PARTIAL_SOLVE_CONTROL cannot be enabled (nonzero
      values are rejected): both read or write through Pardiso's `perm`
      argument, which every handler in _pardiso_ffi.cc passes as null, so
      enabling either dereferences a null pointer.
    - SCHUR_COMPLEMENT_CONTROL cannot be enabled: it needs output buffers
      this package's FFI signatures do not have.
    - TRANSPOSE_SOLVE cannot be set through an overlay at all: use the
      existing `transpose` argument on `solve`/`PardisoSolver.solve` instead,
      which already covers every value this index can meaningfully take for
      the real-valued matrices this package supports.
    """

    USE_DEFAULT_VALUES = 0
    """Whether Pardiso fills in its own defaults (0) or every entry is used as given (nonzero)."""

    FILL_IN_REDUCING_ORDERING = 1
    """Reordering algorithm: 0 minimum degree, 2 METIS nested dissection, 3 parallel nested
    dissection. Default 2."""

    PRECONDITIONED_CGS = 3
    """Krylov subspace iteration and stopping criteria, form 10*L+K with K in {0,1,2}, L>=0.
    Default 0 (off)."""

    USER_PERMUTATION = 4
    """Use of a user-supplied fill-reducing permutation. Default 0. Cannot be enabled here: see
    the class docstring."""

    WRITE_SOLUTION_LOCATION = 5
    """Where the solution is written: 0 to the solution buffer, 1 overwrites the right-hand
    side buffer in place. Default 0."""

    MAX_ITERATIVE_REFINEMENT_STEPS = 7
    """Max iterative refinement iterations for solve: 0 automatic (2 steps), positive an
    explicit cap, negative a cap with extended precision. Default 0."""

    REFINEMENT_TOLERANCE = 8
    """Relative residual tolerance for stopping iterative refinement. Default 0 (Pardiso's own
    default checks apply)."""

    PIVOTING_PERTURBATION = 9
    """Small/zero pivots are perturbed by eps = 10^-iparm[9]. Default 13 for nonsymmetric and
    structurally symmetric matrices, 8 for symmetric ones."""

    SCALING = 10
    """Matrix scaling for diagonal dominance. Default 1 (on) for REAL_NONSYMMETRIC, 0 (off)
    otherwise."""

    TRANSPOSE_SOLVE = 11
    """Which system a solve step solves: 0 Ax=b, 1 conjugate transpose, 2 transpose. Cannot be
    set here: use the `transpose` argument instead, see the class docstring."""

    WEIGHTED_MATCHING = 12
    """Maximum weighted matching for diagonal elements. Default 1 (on) for REAL_NONSYMMETRIC, 0
    (off) otherwise, but this package always defaults it to 0: see the class docstring."""

    REPORT_NONZEROS_IN_FACTORS = 17
    """Report the count of non-zeros in the L and U factors as a diagnostic. Negative enables
    reporting, non-negative disables it. Default -1 (enabled)."""

    REPORT_FACTORIZATION_MFLOPS = 18
    """Report factorization cost in units of 10^6 floating point operations. Negative enables
    reporting, non-negative disables it. Default 0 (disabled)."""

    PIVOTING_STRATEGY = 20
    """Pivoting strategy for symmetric indefinite matrices: 0 pure 1x1, 1 1x1 plus 2x2
    Bunch-Kaufman, 2 and 3 the same without auto-refinement. Default 1."""

    PARALLEL_FACTORIZATION_CONTROL = 23
    """Factorization algorithm variant: 0 classic, 1 two-level, 10 improved two-level
    (nonsymmetric only). Default 0."""

    PARALLEL_SOLVE_CONTROL = 24
    """Parallelization strategy for the solve step: 0 automatic, 1 sequential, 2 matrix
    partitioning. Default 0."""

    MATRIX_CHECKER = 26
    """Validate the ia/ja arrays and column ordering before use. Default 0 (no check)."""

    PRECISION = 27
    """Single (1) or double (0) precision computation. Default 0."""

    PARTIAL_SOLVE_CONTROL = 30
    """Sparse right-hand-side and selective solution output. Default 0. Cannot be enabled here:
    see the class docstring."""

    CNR_THREAD_COUNT = 33
    """OpenMP thread count for conditional numerical reproducibility: 0 auto-determine, positive
    an explicit count. Default 0."""

    INDEXING_STYLE = 34
    """One-based (0) or zero-based (1) array indexing. This package always defaults it to 1.
    Overriding it raises a runtime warning: see the class docstring."""

    SCHUR_COMPLEMENT_CONTROL = 35
    """Compute a Schur complement and in what format. Default 0 (none). Cannot be enabled here:
    see the class docstring."""

    MATRIX_STORAGE_FORMAT = 36
    """Input matrix storage format: 0 CSR, positive a BSR block size, negative a VBSR
    conversion threshold. Default 0."""

    LOW_RANK_UPDATE = 38
    """Enable a low-rank update for a factorization similar to a previous one. Requires
    PARALLEL_FACTORIZATION_CONTROL = 10. Default 0 (off)."""

    INVERSE_DIAGONAL_COMPUTATION = 42
    """Compute the diagonal of the matrix inverse during factorization. Requires
    PARALLEL_FACTORIZATION_CONTROL = 1 and a symmetric matrix type. Default 0 (off). This
    package has no retrieval path for the result, so enabling it has no observable effect here."""

    DIAGONAL_PIVOTING_CALLBACK_CONTROL = 55
    """Enable pivot control and diagonal extraction callbacks (in-core mode only). Default 0
    (off). This package never registers a callback, so enabling it has no observable effect
    here."""

    IN_CORE_MODE = 59
    """In-core (0), automatic (1), or forced out-of-core (2) execution mode. Default 0."""

@jax.tree_util.register_dataclass
@dataclasses.dataclass(frozen=True)
class PardisoDiagnostics:
    """Pardiso's write-back outputs, read from the final iparm array after a call.

    Built by from_iparm, which is plain array indexing with no concretizing
    operation anywhere in it, so it is trace-safe: it produces a correct
    result whether iparm is a concrete array (the ordinary eager case) or
    still traced (when the call producing it is itself wrapped in jax.jit).
    This is why every field here is a JAX scalar array rather than a Python
    int, and why the class is registered as a JAX pytree: a plain dataclass
    of Python ints could not flow through a jitted function as an output.

    Under vmap, every field's array still gains a batch dimension by default
    (jax.vmap's default out_axes=0 broadcasts any output, batched or not),
    but the values differ by case: when a solve genuinely runs once per
    batch element (batching over matrix values), each entry is that
    element's own diagnostics; when one native call already covers the
    whole batch (batching only over right-hand sides), Pardiso reports
    diagnostics once, and every entry along the batch dimension is that same
    value, broadcast rather than recomputed. A caller who wants the compact,
    un-broadcast form in the second case can request it explicitly with
    jax.vmap's own out_axes argument. See the "Batching with vmap" and
    "Diagnostics" sections of the user guide for worked examples.

    Only ever available on a successful call: a Pardiso error raises a
    Python exception instead of returning a value, and a function that
    raises cannot also return diagnostics, on success or partial failure.
    """

    refinement_steps_performed: jax.Array
    """Iterative refinement steps actually run (iparm[6])."""

    perturbed_pivot_count: jax.Array
    """Number of pivots perturbed during factorization (iparm[13])."""

    peak_memory_symbolic_kb: jax.Array
    """Peak memory used during the symbolic (analysis) phase, in KB (iparm[14])."""

    permanent_memory_symbolic_kb: jax.Array
    """Permanent memory retained from the symbolic phase, in KB (iparm[15])."""

    peak_memory_numerical_kb: jax.Array
    """Peak memory used during numerical factorization, in KB (iparm[16])."""

    nonzeros_in_factors: jax.Array
    """Non-zero count in the L and U factors, if REPORT_NONZEROS_IN_FACTORS was set (iparm[17])."""

    factorization_mflops: jax.Array
    """Factorization cost in units of 10^6 FLOPs, if REPORT_FACTORIZATION_MFLOPS was set
    (iparm[18])."""

    cgs_diagnostics: jax.Array
    """Krylov (CG/CGS) iteration diagnostics, if PRECONDITIONED_CGS was set (iparm[19])."""

    positive_eigenvalues: jax.Array
    """Positive eigenvalue count, for symmetric indefinite matrices (iparm[21])."""

    negative_eigenvalues: jax.Array
    """Negative eigenvalue count, for symmetric indefinite matrices (iparm[22])."""

    zero_or_negative_pivot_position: jax.Array
    """Position of the first zero or negative pivot encountered (iparm[29])."""

    min_out_of_core_memory_kb: jax.Array
    """Minimum memory required for out-of-core factorization, in KB (iparm[62])."""

    raw: jax.Array
    """All 64 iparm entries, for anything not decoded into a named field above."""

    @staticmethod
    def from_iparm(iparm: jax.Array) -> PardisoDiagnostics:
        """Decode a final iparm array (shape (..., 64)) into a PardisoDiagnostics."""
        return PardisoDiagnostics(
            refinement_steps_performed=iparm[..., 6],
            perturbed_pivot_count=iparm[..., 13],
            peak_memory_symbolic_kb=iparm[..., 14],
            permanent_memory_symbolic_kb=iparm[..., 15],
            peak_memory_numerical_kb=iparm[..., 16],
            nonzeros_in_factors=iparm[..., 17],
            factorization_mflops=iparm[..., 18],
            cgs_diagnostics=iparm[..., 19],
            positive_eigenvalues=iparm[..., 21],
            negative_eigenvalues=iparm[..., 22],
            zero_or_negative_pivot_position=iparm[..., 29],
            min_out_of_core_memory_kb=iparm[..., 62],
            raw=iparm,
        )

@jax.tree_util.register_pytree_node_class
class FactorizationToken:
    """Handle to a native factorization: a cache id plus a solve counter.

    n_dependent_solutions counts the solutions passed to track. release consumes
    it, so a release is ordered after those solves even inside a jit trace.
    """

    def __init__(self, id, n_dependent_solutions):
        self.id = id
        self.n_dependent_solutions = n_dependent_solutions

    def track(self, *solutions):
        """Return a token whose release is ordered after these solutions."""
        count = self.n_dependent_solutions
        for solution in solutions:
            count = count + jnp.int32(1) + _ordering_witness(solution)
        return FactorizationToken(self.id, count)

    def tree_flatten(self):
        return (self.id, self.n_dependent_solutions), None

    @classmethod
    def tree_unflatten(cls, aux_data, children):
        return cls(*children)