Add (failing) test

patrick-kidger · Aug 7, 2024 · fd30ac0 · fd30ac0
1 parent dcfcbc9
commit fd30ac0
Showing 1 changed file with 7 additions and 1 deletion.
diff --git a/tests/helpers.py b/tests/helpers.py
@@ -317,7 +317,7 @@ def loss(model, x, y):
 
 def square_minus_one(x: Array, args: PyTree):
     """A simple ||x||^2 - 1 function."""
-    return jnp.sum(jnp.square(x)) - 1.0
+    return jnp.sum(jnp.square(jnp.abs(x))) - 1.0
 
 
 #
@@ -463,6 +463,12 @@ def get_weights(model):
     ),
     # Problems with initial value of 0
     (square_minus_one, jnp.array(-1.0), jnp.array(1.0), None),
+    (
+        square_minus_one,
+        jnp.array(-1.0, dtype=jnp.complex128),
+        jnp.array(1.0, dtype=jnp.complex128),
+        None,
+    ),
 )
 
 # ROOT FIND/FIXED POINT PROBLEMS