Merge pull request #20552 from jakevdp:geomspace-complex

PiperOrigin-RevId: 621348396

Author: jax authors

CHANGELOG.md

@@ -12,6 +12,10 @@ Remember to align the itemized text with the first line of an item within a list
   * Added {func}`jax.numpy.trapezoid`, following the addition of this function in
     NumPy 2.0.
 
+* Changes
+  * Complex-valued {func}`jax.numpy.geomspace` now chooses the logarithmic spiral
+    branch consistent with that of NumPy 2.0.
+
 * Deprecations & Removals
   * {func}`jax.tree_map` is deprecated; use `jax.tree.map` instead, or for backward
     compatibility with older JAX versions, use {func}`jax.tree_util.tree_map`.

jax/_src/numpy/lax_numpy.py

@@ -2685,13 +2685,11 @@ def _geomspace(start: ArrayLike, stop: ArrayLike, num: int = 50, endpoint: bool
   util.check_arraylike("geomspace", start, stop)
   start = asarray(start, dtype=computation_dtype)
   stop = asarray(stop, dtype=computation_dtype)
-  # follow the numpy geomspace convention for negative and complex endpoints
-  signflip = 1 - (1 - ufuncs.sign(ufuncs.real(start))) * (1 - ufuncs.sign(ufuncs.real(stop))) // 2
-  signflip = signflip.astype(computation_dtype)
-  res = signflip * logspace(ufuncs.log10(signflip * start),
-                            ufuncs.log10(signflip * stop), num,
-                            endpoint=endpoint, base=10.0,
-                            dtype=computation_dtype, axis=0)
+
+  sign = ufuncs.sign(start)
+  res = sign * logspace(ufuncs.log10(start / sign), ufuncs.log10(stop / sign),
+                        num, endpoint=endpoint, base=10.0,
+                        dtype=computation_dtype, axis=0)
   if axis != 0:
     res = moveaxis(res, 0, axis)
   return lax.convert_element_type(res, dtype)

tests/lax_numpy_test.py

@@ -5068,10 +5068,6 @@ def np_op(start, stop):
   @jax.numpy_rank_promotion('allow')  # This test explicitly exercises implicit rank promotion.
   def testGeomspace(self, start_shape, stop_shape, num,
                     endpoint, dtype, axis):
-    # TODO(jakevdp): remove once https://github.com/numpy/numpy/issues/26195 is fixed.
-    if jtu.numpy_version() >= (2, 0, 0) and dtypes.issubdtype(dtype, jnp.complexfloating):
-      self.skipTest("NumPy 2.0.0rc1 geomspace is broken for complex dtypes.")
-
     rng = jtu.rand_default(self.rng())
     # relax default tolerances slightly
     tol = {dtypes.bfloat16: 2e-2, np.float16: 4e-3, np.float32: 2e-3,
@@ -5101,8 +5097,11 @@ def np_op(start, stop):
         start, stop, num, endpoint=endpoint,
         dtype=dtype if dtype != jnp.bfloat16 else np.float32,
         axis=axis).astype(dtype)
-    self._CheckAgainstNumpy(np_op, jnp_op, args_maker,
-                            check_dtypes=False, tol=tol)
+
+    # JAX follows NumPy 2.0 semantics for complex geomspace.
+    if not (jtu.numpy_version() < (2, 0, 0) and dtypes.issubdtype(dtype, jnp.complexfloating)):
+      self._CheckAgainstNumpy(np_op, jnp_op, args_maker,
+                              check_dtypes=False, tol=tol)
     if dtype in (inexact_dtypes + [None,]):
       self._CompileAndCheck(jnp_op, args_maker,
                             check_dtypes=False, atol=tol, rtol=tol)