Merge pull request #1336 from vilchy/fix-complex-sqrt

Fix sqrt for complex numbers when x.im == 0
google-research · Nov 14, 2023 · be6befd · be6befd
2 parents 3cbde4c + 8fb4118
commit be6befd
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 1 deletion.
diff --git a/lib/complex.dx b/lib/complex.dx
@@ -58,7 +58,8 @@ def complex_pow(base:Complex, power:Complex) -> Complex = complex_exp(power * co
 
 def complex_sqrt(x:Complex) -> Complex =
   m = complex_mag x
-  Complex (sqrt ((x.re + m) / 2.0)) (sign x.im * sqrt ((m - x.re) / 2.0))
+  sgn = if x.im >= 0 then 1.0 else -1.0
+  Complex (sqrt ((x.re + m) / 2.0)) (sgn * sqrt ((m - x.re) / 2.0))
 
 def complex_sin( x:Complex) -> Complex = Complex(sin  x.re * cosh x.im, cos  x.re * sinh x.im)
 def complex_sinh(x:Complex) -> Complex = Complex(sinh x.re *  cos x.im, cosh x.re * sin  x.im)

diff --git a/tests/complex-tests.dx b/tests/complex-tests.dx
@@ -34,6 +34,8 @@ b = Complex (-1.1) 1.3
 > True
 :p sqrt (sq a) ~~ a
 > True
+:p sqrt (Complex (-1.0) 0.0) ~~ (Complex 0.0 1.0)
+> True
 :p log ((Complex 1.0 0.0) + a) ~~ log1p a
 > True
 :p sin (-a) ~~ (-(sin a))