Merge pull request #33837 from jakevdp:arange-complex-dep

PiperOrigin-RevId: 842354532

Author: Google-ML-Automation

CHANGELOG.md

@@ -21,6 +21,8 @@ When releasing, please add the new-release-boilerplate to docs/pallas/CHANGELOG.
     Please use `jax.lax.pcast(..., to='varying')` as the replacement.
   * `with mesh:` context manager has been deprecated.
     Please use `with jax.set_mesh(mesh):` instead.
+  * Complex arguments passed to {func}`jax.numpy.arange` now result in a
+    deprecation warning, because the output is poorly-defined.
 
 * Changes:
   * jax's `Tracer` no longer inherits from `jax.Array` at runtime. However,

jax/_src/deprecations.py

@@ -128,7 +128,8 @@ def warn(deprecation_id: str, message: str, stacklevel: int) -> None:
 register('jax-lax-dot-positional-args')
 register('jax-lib-module')
 register('jax-nn-one-hot-float-input')
-register("jax-numpy-astype-complex-to-real")
+register('jax-numpy-arange-complex')
+register('jax-numpy-astype-complex-to-real')
 register('jax-numpy-clip-args')
 register('jax-scipy-special-sph-harm')
 register('safer-randint-config')

jax/_src/numpy/lax_numpy.py

@@ -5983,6 +5983,14 @@ def _arange(start: ArrayLike | DimSize, stop: ArrayLike | DimSize | None = None,
   dtype = dtypes.jax_dtype(dtype)
 
   if iscomplexobj(start) or iscomplexobj(stop) or iscomplexobj(step):
+    deprecations.warn(
+      "jax-numpy-arange-complex",
+      (
+        "Passing complex start/stop/step to jnp.arange is deprecated;"
+        " in the future this will result in a ValueError."
+      ),
+      stacklevel=3
+    )
     # Complex arange is poorly defined; fall back to NumPy here.
     # TODO(jakevdp): deprecate the complex case.
     return array(np.arange(start, stop, step, dtype=dtype), device=out_sharding)

tests/lax_numpy_test.py

@@ -46,6 +46,7 @@
 from jax._src import array
 from jax._src import config
 from jax._src import core
+from jax._src import deprecations
 from jax._src import dtypes
 from jax._src import test_util as jtu
 from jax._src.lax import lax as lax_internal
@@ -4892,18 +4893,23 @@ def testArangeRandomValues(self, dtype, iteration):
     np_result = np.arange(start, stop, dtype=dtype)
     self.assertAllClose(jax_result, np_result)
 
-  def testArangeComplex(self):
-    test_cases = [
+  @parameterized.parameters(
       (1+2j, 5+3j),
       (0+0j, 5+0j),
       (1.0+0j, 5.0+0j),
       (0, 5, 1+1j),
-    ]
-    for args in test_cases:
-      with self.subTest(args=args):
+  )
+  def testArangeComplex(self, *args):
+    dep_id = "jax-numpy-arange-complex"
+    msg = "Passing complex start/stop/step to jnp.arange is deprecated"
+    if deprecations.is_accelerated(dep_id):
+      with self.assertRaisesRegex(ValueError, msg):
+        jax_result = jnp.arange(*args)
+    else:
+      with self.assertWarnsRegex(DeprecationWarning, msg):
         jax_result = jnp.arange(*args)
-        np_result = np.arange(*args)
-        self.assertArraysEqual(jax_result, np_result)
+      np_result = np.arange(*args)
+      self.assertArraysEqual(jax_result, np_result)
 
   def testIssue830(self):
     a = jnp.arange(4, dtype=jnp.complex64)