Merge pull request #33602 from yashwantbezawada:fix-arange-constant-folding

PiperOrigin-RevId: 841846305

Author: Google-ML-Automation

jax/_src/numpy/lax_numpy.py

@@ -5989,9 +5989,23 @@ def _arange(start: ArrayLike | DimSize, stop: ArrayLike | DimSize | None = None,
       start = ceil_(start).astype(int)
     return lax.broadcasted_iota(dtype, (start,), 0, out_sharding=out_sharding)  # type: ignore[arg-type]
   else:
-    if step is None and start == 0 and stop is not None:
-      return lax.broadcasted_iota(dtype, (np.ceil(stop).astype(int),), 0,
-                                  out_sharding=out_sharding)
+    if step is None and stop is not None:
+      # Skip optimization if start or stop is complex (ceil doesn't support complex)
+      start_dtype = _dtype(start)
+      stop_dtype = _dtype(stop)
+      if (dtypes.issubdtype(start_dtype, np.complexfloating) or
+          dtypes.issubdtype(stop_dtype, np.complexfloating)):
+        return array(np.arange(start, stop=stop, step=step, dtype=dtype),
+                     device=out_sharding)
+      # Use iota + offset instead of creating a constant array
+      size = int(np.ceil(stop - start))
+      if size <= 0:
+        return array([], dtype=dtype, device=out_sharding)
+      result = lax.broadcasted_iota(dtype, (size,), 0, out_sharding=out_sharding)
+      if start != 0:
+        # Add offset if start is non-zero
+        result = lax.add(result, lax.convert_element_type(start, dtype))
+      return result
     return array(np.arange(start, stop=stop, step=step, dtype=dtype),
                  device=out_sharding)
 

tests/lax_numpy_test.py

@@ -4869,6 +4869,42 @@ def testArangeJaxpr(self, args, specify_device):
     self.assertEqual(len(jaxpr.jaxpr.eqns), num_eqs)
     self.assertEqual(jaxpr.jaxpr.eqns[0].primitive, lax.iota_p)
 
+  @jtu.sample_product(specify_device=[True, False])
+  def testArangeJaxprNonZeroStart(self, specify_device):
+    device = jax.devices()[-1] if specify_device else None
+    jaxpr = jax.make_jaxpr(lambda: jnp.arange(1, 5, device=device))()
+    # Non-zero start should produce iota + add (+ device_put if device specified)
+    num_eqs = 3 if device is not None else 2
+    self.assertEqual(len(jaxpr.jaxpr.eqns), num_eqs)
+    self.assertEqual(jaxpr.jaxpr.eqns[0].primitive, lax.iota_p)
+    self.assertEqual(jaxpr.jaxpr.eqns[1].primitive, lax.add_p)
+
+  @jtu.sample_product(
+      dtype=[np.int32, np.float32],
+      iteration=range(10)
+  )
+  def testArangeRandomValues(self, dtype, iteration):
+    del iteration  # not needed: each test case gets its own random seed.
+    rng = jtu.rand_default(self.rng())
+    start = rng((), dtype)
+    stop = rng((), dtype)
+    jax_result = jnp.arange(start, stop, dtype=dtype)
+    np_result = np.arange(start, stop, dtype=dtype)
+    self.assertAllClose(jax_result, np_result)
+
+  def testArangeComplex(self):
+    test_cases = [
+      (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):
+        jax_result = jnp.arange(*args)
+        np_result = np.arange(*args)
+        self.assertArraysEqual(jax_result, np_result)
+
   def testIssue830(self):
     a = jnp.arange(4, dtype=jnp.complex64)
     self.assertEqual(a.dtype, jnp.complex64)