--- Final Solver (used for performance test) ---
import numpy as np
import warnings
from scipy.linalg._matfuncs_schur_sqrtm import recursive_schur_sqrtm


class Solver:
    def solve(self, problem, **kwargs):
        A = np.asarray(problem["matrix"], dtype=np.complex128)

        n = A.shape[0]
        if n == 1:
            return {"sqrtm": {"X": [[np.sqrt(A[0, 0]).tolist()]]}}

        if n == 2:
            X = _sqrtm_2x2(A)
            return {"sqrtm": {"X": X.tolist()}}

        X, _, _, _ = recursive_schur_sqrtm(A)
        return {"sqrtm": {"X": X.tolist()}}


def _sqrtm_2x2(A):
    a, b = A[0, 0], A[0, 1]
    c, d = A[1, 0], A[1, 1]
    det = a * d - b * c
    s = np.sqrt(det)
    t = np.sqrt(a + d + 2.0 * s)
    if abs(t) < 1e-300:
        X, _, _, _ = recursive_schur_sqrtm(A)
        return X
    return (A + s * np.eye(2, dtype=np.complex128)) / t
--- 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: 32.2534s
Total Solver Time:   11.6387s
Raw Speedup:         2.7712 x
Final Reward (Score): 2.7712
---------------------------
..

=============================== warnings summary ===============================
test_outputs.py: 1100 warnings
  /tests/evaluator.py:79: DeprecationWarning: The `disp` argument is deprecated and will be removed in SciPy 1.18.0.
    X, _ = scipy.linalg.sqrtm(

-- Docs: https://docs.pytest.org/en/stable/how-to/capture-warnings.html
==================================== 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, 1100 warnings in 639.33s (0:10:39) =================
