--- Final Solver (used for performance test) ---
import warnings

import numpy as np
import scipy.linalg
import threadpoolctl


class Solver:
    def __init__(self):
        # The benchmark environment uses OpenBLAS with many threads. For the
        # matrix sizes used here, multi-threaded BLAS/LAPACK causes severe
        # overhead. Limiting BLAS to one thread gives a large speedup,
        # especially for n >= 100. This is done in __init__ so it does not
        # count toward solve() runtime.
        self._thread_limit = threadpoolctl.threadpool_limits(
            limits=1, user_api="blas"
        )
        # Avoid runtime warnings from scipy.linalg.sqrtm in the output stream.
        warnings.simplefilter("ignore", category=Warning)

    def solve(self, problem, **kwargs):
        A = np.asarray(problem["matrix"], dtype=complex)
        if A.ndim != 2 or A.shape[0] != A.shape[1]:
            return {"sqrtm": {"X": []}}

        n = A.shape[0]
        if n == 0:
            X = np.empty((0, 0), dtype=complex)
        elif n == 1:
            X = np.array([[np.sqrt(A[0, 0])]], dtype=complex)
        elif n == 2:
            X = _sqrtm2(A)
            if X is None:
                try:
                    X = scipy.linalg.sqrtm(A)
                except Exception:
                    return {"sqrtm": {"X": []}}
        else:
            try:
                X = scipy.linalg.sqrtm(A)
            except Exception:
                return {"sqrtm": {"X": []}}

        X = np.asarray(X, dtype=complex)
        if X.shape != (n, n):
            X = X.reshape((n, n))
        return {"sqrtm": {"X": X.tolist()}}


def _sqrtm2(A):
    """
    Principal square root for a 2x2 complex matrix.

    Returns None if the denominator is too small, in which case the caller can
    fall back to scipy.linalg.sqrtm.
    """
    a = A[0, 0]
    b = A[0, 1]
    c = A[1, 0]
    d = A[1, 1]

    tr = a + d
    det = a * d - b * c
    disc = tr * tr - 4.0 * det
    sq = np.sqrt(disc)

    l1 = (tr + sq) / 2.0
    l2 = (tr - sq) / 2.0
    m1 = np.sqrt(l1)
    m2 = np.sqrt(l2)
    s = m1 * m2

    denom = np.sqrt(tr + 2.0 * s)
    if not np.isfinite(denom) or np.abs(denom) < 1e-14:
        return None

    return (A + s * np.eye(2, dtype=complex)) / denom
--- End Solver ---
Running performance test...
============================= test session starts ==============================
platform linux -- Python 3.12.13, pytest-9.0.2, pluggy-1.6.0
rootdir: /tests
plugins: anyio-4.13.0, jaxtyping-0.3.9
collected 3 items

../tests/test_outputs.py .
--- Performance Summary ---
Validity: True
Total Baseline Time: 3.8978s
Total Solver Time:   3.7923s
Raw Speedup:         1.0278 x
Final Reward (Score): 1.0278
---------------------------
..

==================================== PASSES ====================================
=========================== short test summary info ============================
PASSED ../tests/test_outputs.py::test_solver_exists
PASSED ../tests/test_outputs.py::test_solver_validity
PASSED ../tests/test_outputs.py::test_solver_speedup
========================= 3 passed in 87.79s (0:01:27) =========================
