Repository: google/jax
Issue: Inconsistent behaviour of arange

I have installed JAX on Colab as in the QuickStart notebook. Then I run this instruction:

```
import jax.numpy as np
print(np.arange(0.0, 1.0, 0.1))
```

And I get an array with a single value [0], which is different from:

```
import numpy as np
print(np.arange(0.0, 1.0, 0.1))
```

Where I (correctly) get an array with 11 values 0, 0.1, ..., 0.9, 1.0. Is this intended?

Changed files:
jax/numpy/lax_numpy.py, setup.py, tests/lax_numpy_test.py

diff --git a/jax/numpy/lax_numpy.py b/jax/numpy/lax_numpy.py
index 6329f9b31..2cb34e79c 100644
--- a/jax/numpy/lax_numpy.py
+++ b/jax/numpy/lax_numpy.py
@@ -834,29 +834,17 @@ def eye(N, M=None, k=None, dtype=onp.dtype("float64")):
 
 @_wraps(onp.arange)
 def arange(*args, **kwargs):
-  nargs = len(args)
-  start, step = 0, 1
-  dtype = kwargs.pop("dtype", None)
-  if kwargs:
-    raise TypeError("arange only accepts 'dtype' kwarg, got {}".format(kwargs))
-  if nargs == 0:
-    raise TypeError("Required argument 'start' (pos 1) not found")  # same as numpy error
-  elif nargs == 1:
-    stop, = args
-    dtype = dtype or _dtype(stop)
-    return lax.iota(dtype, stop)  # avoids materializing
-  elif nargs == 2:
-    start, stop = args
-    dtype = dtype or onp.result_type(start, stop)
-  elif nargs == 3:
-    start, stop, step = args
-    dtype = dtype or onp.result_type(start, stop, step)
-  elif nargs == 4:
-    start, stop, step, dtype = args
-    dtype = dtype or onp.result_type(start, stop, step)
-
-  size = (stop - start - 1) // step + 1
-  return start + step * lax.iota(dtype, size)
+  # attempt to generate a lazy IotaConstant, otherwise fall back to raw numpy
+  # TODO(mattjj): add tests for this function, then re-enable
+  # dtype = kwargs.pop("dtype", None)
+  # if not args:
+  #   raise TypeError("Required argument 'start' (pos 1) not found")  # same as numpy error
+  # elif len(args) == 1 and not kwargs:
+  #   stop, = args
+  #   dtype = dtype or _dtype(stop)
+  #   if onp.issubdtype(dtype, onp.integer):
+  #     return lax.iota(dtype, stop)  # avoids materializing
+  return onp.arange(*args, **kwargs)
 
 
 @_wraps(onp.repeat)
diff --git a/setup.py b/setup.py
index 582aed3f5..e3d98c954 100644
--- a/setup.py
+++ b/setup.py
@@ -16,7 +16,7 @@ from setuptools import setup, find_packages
 
 setup(
     name='jax',
-    version='0.1.12',
+    version='0.1.13',
     description='Differentiate, compile, and transform Numpy code.',
     author='JAX team',
     author_email='jax-dev@google.com',
diff --git a/tests/lax_numpy_test.py b/tests/lax_numpy_test.py
index 52aa98996..e7b80ef0d 100644
--- a/tests/lax_numpy_test.py
+++ b/tests/lax_numpy_test.py
@@ -845,6 +845,13 @@ class LaxBackedNumpyTests(jtu.JaxTestCase):
     ans = onp.mean(x)
     self.assertAllClose(ans, onp.array([1./3, 1./3, 1./3]), check_dtypes=False)
 
+  # TODO(mattjj): more exhaustive arange tests
+  def testArangeOnFloats(self):
+    # from https://github.com/google/jax/issues/145
+    expected = onp.arange(0.0, 1.0, 0.1)
+    ans = lnp.arange(0.0, 1.0, 0.1)
+    self.assertAllClose(expected, ans, check_dtypes=True)
+
 
 if __name__ == "__main__":
   absltest.main()