Repository: google/jax
Issue: DeviceArray.reshape(*shape) does not work

NumPy supports both `.reshape(*shape)` and `.reshape(shape)`, but only the later works with jax.numpy:
```
>>> import jax.numpy as jnp
>>> jnp.arange(12).reshape(3, 4)
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-27-144a56eb5cb3> in <module>()
----> 1 jnp.arange(12).reshape(3, 4)

/usr/local/lib/python3.6/dist-packages/jax/numpy/lax_numpy.py in _reshape(a, newshape, order)
    638 def _reshape(a, newshape, order="C"):
    639   dummy_val = onp.broadcast_to(0, shape(a))  # zero strides
--> 640   computed_newshape = onp.reshape(dummy_val, newshape).shape
    641 
    642   if order == "C":

/usr/local/lib/python3.6/dist-packages/numpy/core/fromnumeric.py in reshape(a, newshape, order)
    290            [5, 6]])
    291     """
--> 292     return _wrapfunc(a, 'reshape', newshape, order=order)
    293 
    294 

/usr/local/lib/python3.6/dist-packages/numpy/core/fromnumeric.py in _wrapfunc(obj, method, *args, **kwds)
     54 def _wrapfunc(obj, method, *args, **kwds):
     55     try:
---> 56         return getattr(obj, method)(*args, **kwds)
     57 
     58     # An AttributeError occurs if the object does not have

ValueError: cannot reshape array of size 12 into shape (3,)
```

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 6afc0fbf3..b1e59e40b 100644
--- a/jax/numpy/lax_numpy.py
+++ b/jax/numpy/lax_numpy.py
@@ -649,6 +649,20 @@ def _reshape(a, newshape, order="C"):
   else:
     raise ValueError("Unexpected value for 'order' argument: {}.".format(order))
 
+def _reshape_method(a, *newshape, **kwargs):
+  order = kwargs.pop("order", "C")
+  if len(kwargs) == 1:
+    invalid_kwarg, = kwargs
+    msg = "'{}' is an invalid keyword argument for this function"
+    raise TypeError(msg.format(invalid_kwarg))  # same as NumPy error
+  elif kwargs:
+    invalid_kwargs = "'{}'".format("'".join(kwargs))
+    msg = "{} are invalid keyword arguments for this function"
+    raise TypeError(msg.format(invalid_kwargs))  # different from NumPy error
+  if len(newshape) == 1 and not isinstance(newshape[0], int):
+    newshape = newshape[0]
+  return _reshape(a, newshape, order=order)
+
 
 @_wraps(onp.ravel)
 def ravel(a, order="C"):
@@ -2353,7 +2367,7 @@ for operator_name, function in _operators.items():
 # Forward methods and properties using core.aval_method and core.aval_property:
 for method_name in _nondiff_methods + _diff_methods:
   setattr(ShapedArray, method_name, core.aval_method(globals()[method_name]))
-setattr(ShapedArray, "reshape", core.aval_method(_reshape))
+setattr(ShapedArray, "reshape", core.aval_method(_reshape_method))
 setattr(ShapedArray, "flatten", core.aval_method(ravel))
 setattr(ShapedArray, "T", core.aval_property(transpose))
 setattr(ShapedArray, "real", core.aval_property(real))
@@ -2367,7 +2381,7 @@ for operator_name, function in _operators.items():
   setattr(DeviceArray, "__{}__".format(operator_name), function)
 for method_name in _nondiff_methods + _diff_methods:
   setattr(DeviceArray, method_name, globals()[method_name])
-setattr(DeviceArray, "reshape", _reshape)
+setattr(DeviceArray, "reshape", _reshape_method)
 setattr(DeviceArray, "flatten", ravel)
 setattr(DeviceArray, "T", property(transpose))
 setattr(DeviceArray, "real", property(real))
diff --git a/tests/lax_numpy_test.py b/tests/lax_numpy_test.py
index f1d692deb..e04a16b3d 100644
--- a/tests/lax_numpy_test.py
+++ b/tests/lax_numpy_test.py
@@ -975,6 +975,25 @@ class LaxBackedNumpyTests(jtu.JaxTestCase):
     self._CheckAgainstNumpy(onp_fun, lnp_fun, args_maker, check_dtypes=True)
     self._CompileAndCheck(lnp_fun, args_maker, check_dtypes=True)
 
+  @parameterized.named_parameters(jtu.cases_from_list(
+      {"testcase_name": "_inshape={}_outshape={}".format(
+          jtu.format_shape_dtype_string(arg_shape, dtype),
+          jtu.format_shape_dtype_string(out_shape, dtype)),
+       "arg_shape": arg_shape, "out_shape": out_shape, "dtype": dtype,
+       "rng": jtu.rand_default()}
+      for dtype in default_dtypes
+      for arg_shape, out_shape in [
+          ((7, 0), (0, 42, 101)),
+          ((2, 1, 4), (-1,)),
+          ((2, 2, 4), (2, 8))
+      ]))
+  def testReshapeMethod(self, arg_shape, out_shape, dtype, rng):
+    onp_fun = lambda x: onp.reshape(x, out_shape)
+    lnp_fun = lambda x: x.reshape(*out_shape)
+    args_maker = lambda: [rng(arg_shape, dtype)]
+    self._CheckAgainstNumpy(onp_fun, lnp_fun, args_maker, check_dtypes=True)
+    self._CompileAndCheck(lnp_fun, args_maker, check_dtypes=True)
+
   @parameterized.named_parameters(jtu.cases_from_list(
       {"testcase_name": "_inshape={}_expanddim={}".format(
           jtu.format_shape_dtype_string(arg_shape, dtype), dim),
@@ -1522,6 +1541,9 @@ class LaxBackedNumpyTests(jtu.JaxTestCase):
     assert lnp.allclose(lnp.eye(5000), onp.eye(5000))
     self.assertEqual(0, onp.sum(lnp.eye(1050) - onp.eye(1050)))
 
+  def testIssue746(self):
+    lnp.arange(12).reshape(3, 4)  # doesn't crash
+
 
 if __name__ == "__main__":
   absltest.main()