Repository: google/jax
Issue: atleast_{1,2,3}d fail for scalars

`numpy.atleast_{1,2,3}d` add dimensions for scalars up to the necessary amount, but JAX implementations fail, since `jax.numpy.array` doesn't return 0-dim arrays for scalars (#121).

Examples:
```
In [3]: onp.atleast_1d(1)
Out[3]: array([1])

In [4]: np.atleast_1d(1)
---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
<ipython-input-4-84084c6642da> in <module>
----> 1 np.atleast_1d(1)

~/src/jax/jax/numpy/lax_numpy.py in atleast_1d(*arys)
   1182     arr = array(arys[0])
-> 1183     return arr if arr.ndim >= 1 else arr.reshape(-1)
   1184   else:
   1185     return [atleast_1d(arr) for arr in arys]

AttributeError: 'int' object has no attribute 'ndim'

In [5]: np.atleast_1d(np.int64(1))
Out[5]: array([1])
```

Numpy dtypes (like `np.int64`) have a `ndim` attribute, which allows them to work properly. The testing suite doesn't catch this because it seems to use Numpy dtypes instead of `int` or `float` literals. 

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

diff --git a/jax/numpy/lax_numpy.py b/jax/numpy/lax_numpy.py
index 06e23d032..39f1423cd 100644
--- a/jax/numpy/lax_numpy.py
+++ b/jax/numpy/lax_numpy.py
@@ -1180,7 +1180,7 @@ def column_stack(tup):
 def atleast_1d(*arys):
   if len(arys) == 1:
     arr = array(arys[0])
-    return arr if arr.ndim >= 1 else arr.reshape(-1)
+    return arr if ndim(arr) >= 1 else reshape(arr, -1)
   else:
     return [atleast_1d(arr) for arr in arys]
 
@@ -1189,7 +1189,7 @@ def atleast_1d(*arys):
 def atleast_2d(*arys):
   if len(arys) == 1:
     arr = array(arys[0])
-    return arr if arr.ndim >= 2 else arr.reshape((1, -1))
+    return arr if ndim(arr) >= 2 else reshape(arr, (1, -1))
   else:
     return [atleast_2d(arr) for arr in arys]
 
@@ -1199,9 +1199,9 @@ def atleast_3d(*arys):
   if len(arys) == 1:
     arr = array(arys[0])
     if ndim(arr) <= 1:
-      arr = arr.reshape((1, -1, 1))
+      arr = reshape(arr, (1, -1, 1))
     elif ndim(arr) == 2:
-      arr = arr.reshape(shape(arr) + (1,))
+      arr = reshape(arr, shape(arr) + (1,))
     return arr
   else:
     return [atleast_3d(arr) for arr in arys]
diff --git a/tests/lax_numpy_test.py b/tests/lax_numpy_test.py
index 63dfafd8d..393b856cc 100644
--- a/tests/lax_numpy_test.py
+++ b/tests/lax_numpy_test.py
@@ -1373,6 +1373,22 @@ class LaxBackedNumpyTests(jtu.JaxTestCase):
     expected = onp.reshape(a, (3, 2), order='F')
     self.assertAllClose(ans, expected, check_dtypes=True)
 
+
+  @parameterized.named_parameters(jtu.cases_from_list(
+      {"testcase_name": "_op={}_dtype={}".format(
+          op, {bool: "bool", int: "int", float: "float"}[dtype]),
+       "dtype": dtype, "op": op}
+      for dtype in [int, float, bool]
+      for op in ["atleast_1d", "atleast_2d", "atleast_3d"]))
+  def testAtLeastNdLiterals(self, dtype, op):
+    # Fixes: https://github.com/google/jax/issues/634
+    onp_fun = lambda arg: getattr(onp, op)(arg)
+    lnp_fun = lambda arg: getattr(lnp, op)(arg)
+    args_maker = lambda: [dtype(2)]
+    self._CheckAgainstNumpy(onp_fun, lnp_fun, args_maker, check_dtypes=True)
+    self._CompileAndCheck(lnp_fun, args_maker, check_dtypes=True)
+
+
   def testLongLong(self):
     # TODO(phawkins): enable after a Jaxlib update.
     return SkipTest("Test disabled until jaxlib 0.1.13 is released.")