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

@numba.njit
def _sqrtm_2x2(a, b, c, d):
    s = a + d
    t = np.sqrt((a - d)**2 + 4.0 * b * c)
    l1 = (s + t) / 2.0
    l2 = (s - t) / 2.0
    sqrt_l1 = np.sqrt(l1)
    sqrt_l2 = np.sqrt(l2)
    if abs(t) > 1e-14:
        alpha = (sqrt_l1 + sqrt_l2) / 2.0
        beta = (sqrt_l1 - sqrt_l2) / t
    else:
        alpha = sqrt_l1
        beta = 1.0 / (2.0 * alpha) if abs(alpha) > 1e-14 else 0.0
    return (
        alpha + beta * (a - s / 2.0),
        beta * b,
        beta * c,
        alpha + beta * (d - s / 2.0),
    )

class Solver:
    def __init__(self):
        pass

    def solve(self, problem, **kwargs) -> Any:
        A = problem["matrix"]
        
        if not isinstance(A, np.ndarray):
            A = np.asarray(A, dtype=np.complex128)
        elif A.dtype != np.complex128:
            A = A.astype(np.complex128)
        
        n = A.shape[0]
        
        if n == 1:
            return {"sqrtm": {"X": [[complex(np.sqrt(A[0, 0]))]]}}
        
        if n == 2:
            X00, X01, X10, X11 = _sqrtm_2x2(A[0, 0], A[0, 1], A[1, 0], A[1, 1])
            return {
                "sqrtm": {
                    "X": [
                        [complex(X00), complex(X01)],
                        [complex(X10), complex(X11)],
                    ]
                }
            }
        
        # Cheap Hermitian check (much faster than np.allclose)
        if np.max(np.abs(A - A.conj().T)) < 1e-12:
            try:
                w, v = scipy.linalg.eigh(A)
                sqrt_w = np.sqrt(w.astype(np.complex128))
                X = v @ np.diag(sqrt_w) @ v.conj().T
                return {"sqrtm": {"X": X.tolist()}}
            except Exception:
                pass
        
        # General case: direct compiled call
        res, isIllconditioned, isSingular, info = recursive_schur_sqrtm(A)
        if info >= 0:
            return {"sqrtm": {"X": res.tolist()}}
        
        try:
            X = scipy.linalg.sqrtm(A)
            return {"sqrtm": {"X": X.tolist()}}
        except Exception:
            return {"sqrtm": {"X": []}}
--- 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: 50.2652s
Total Solver Time:   22.0613s
Raw Speedup:         2.2784 x
Final Reward (Score): 2.2784
---------------------------
..

=============================== 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 1042.50s (0:17:22) =================
