PrimeIntellect/SWE-rebench-V2 huggingface dataset:
Repository: mpmath/mpmath (python)

Issue:
sqrt returns incorrect results at complex infinities
As in the title.

Reproducer (mpmath version 1.4.0a1.dev43+g0d287bf):
```python
>>> mpmath.mp.sqrt(mpmath.mp.mpc(mpmath.mp.inf, mpmath.mp.inf))  # expected (+∞,+∞)
mpc(real='+inf', imag='nan')
>>> mpmath.mp.sqrt(mpmath.mp.mpc(-mpmath.mp.inf, mpmath.mp.inf))  # expected (+∞,+∞)
mpc(real='nan', imag='+inf')
>>> mpmath.mp.sqrt(mpmath.mp.mpc(mpmath.mp.inf, -mpmath.mp.inf))  # expected (+∞,-∞)
mpc(real='+inf', imag='nan')
>>> mpmath.mp.sqrt(mpmath.mp.mpc(-mpmath.mp.inf, -mpmath.mp.inf))  # expected (+∞,-∞)
mpc(real='nan', imag='-inf')
```

`mpmath.fp` gives expected results as defined by https://en.cppreference.com/w/cpp/numeric/complex/sqrt.

Pull request:
sqrt(z): special case for infinite z.imag
See https://en.cppreference.com/w/c/numeric/complex/csqrt

Close #776

Patch:
diff --git a/mpmath/libmp/libmpc.py b/mpmath/libmp/libmpc.py
index bc30a94..cfa7df7 100644
--- a/mpmath/libmp/libmpc.py
+++ b/mpmath/libmp/libmpc.py
@@ -295,6 +295,8 @@ def mpc_sqrt(z, prec, rnd=round_fast):
         else:
             re = mpf_sqrt(a, prec, rnd)
             return (re, fzero)
+    if b in (finf, fninf):
+        return (finf, b)
     wp = prec+20
     if not a[0]:                               # case a positive
         t  = mpf_add(mpc_abs((a, b), wp), a, wp)  # t = abs(a+bi) + a


Test patch:
diff --git a/mpmath/tests/test_functions.py b/mpmath/tests/test_functions.py
index 78f1535..0b46a5a 100644
--- a/mpmath/tests/test_functions.py
+++ b/mpmath/tests/test_functions.py
@@ -142,6 +142,14 @@ def test_float_sqrt():
     pytest.raises(ComplexResult, lambda: mp2.mpf(-1)**0.5)
     pytest.raises(ComplexResult, lambda: mp2.mpf(-1)**mp2.mpf(0.5))
 
+def test_sqrt_special():
+    assert sqrt(mpc(+inf, +inf)) == mpc(inf, +inf)
+    assert sqrt(mpc(-inf, +inf)) == mpc(inf, +inf)
+    assert sqrt(mpc( nan, +inf)) == mpc(inf, +inf)
+    assert sqrt(mpc(+inf, -inf)) == mpc(inf, -inf)
+    assert sqrt(mpc(-inf, -inf)) == mpc(inf, -inf)
+    assert sqrt(mpc( nan, -inf)) == mpc(inf, -inf)
+
 def test_hypot():
     assert hypot(0, 0) == 0
     assert hypot(0, 0.33) == mpf(0.33)