Repository: google/jax

Issue:
Fortran reshape order runs but gives wrong behaviour
Hi there,

Big fan of Jax, thanks for making this!

I believe there is a problem with specifying the reshaping order:

```python
a = np.arange(6) + 1
np.reshape(a, (3, 2), order='F')
```

Should give:

```
array([[1, 4],
       [2, 5],
       [3, 6]])
```

But instead gives:

```
array([[1, 2],
       [3, 4],
       [5, 6]], dtype=int32)
```

Basically, it is ignored and always does column-major reordering. If unsupported, it would be nice for it to just fail rather than silently carry on.
Fortran reshape order runs but gives wrong behaviour
Hi there,

Big fan of Jax, thanks for making this!

I believe there is a problem with specifying the reshaping order:

```python
a = np.arange(6) + 1
np.reshape(a, (3, 2), order='F')
```

Should give:

```
array([[1, 4],
       [2, 5],
       [3, 6]])
```

But instead gives:

```
array([[1, 2],
       [3, 4],
       [5, 6]], dtype=int32)
```

Basically, it is ignored and always does column-major reordering. If unsupported, it would be nice for it to just fail rather than silently carry on.


Discussion:

Resolving patch:
diff --git a/jax/numpy/lax_numpy.py b/jax/numpy/lax_numpy.py
--- a/jax/numpy/lax_numpy.py
+++ b/jax/numpy/lax_numpy.py
@@ -597,19 +597,19 @@ def angle(x):
 
 @_wraps(onp.reshape)
 def reshape(a, newshape, order="C"):  # pylint: disable=missing-docstring
-  if order == "C" or order is None:
-    dims = None
+  dummy_val = onp.broadcast_to(0, shape(a))  # zero strides
+  computed_newshape = onp.reshape(dummy_val, newshape).shape
+
+  if order == "C":
+    return lax.reshape(a, computed_newshape, None)
   elif order == "F":
     dims = onp.arange(ndim(a))[::-1]
+    return lax.reshape(a, computed_newshape[::-1], dims).T
   elif order == "A":
     raise NotImplementedError("np.reshape order=A is not implemented.")
   else:
     raise ValueError("Unexpected value for 'order' argument: {}.".format(order))
 
-  dummy_val = onp.broadcast_to(0, shape(a))  # zero strides
-  computed_newshape = onp.reshape(dummy_val, newshape).shape
-  return lax.reshape(a, computed_newshape, dims)
-
 
 @_wraps(onp.ravel)
 def ravel(a, order="C"):


Test patch:
diff --git a/tests/lax_numpy_test.py b/tests/lax_numpy_test.py
--- a/tests/lax_numpy_test.py
+++ b/tests/lax_numpy_test.py
@@ -859,12 +859,14 @@ def fn(module, axis):
     self._CompileAndCheck(lnp_fun, args_maker, check_dtypes=True)
 
   @parameterized.named_parameters(jtu.cases_from_list(
-      {"testcase_name": "_inshape={}_outshape={}".format(
+      {"testcase_name": "_inshape={}_outshape={}_order={}".format(
           jtu.format_shape_dtype_string(arg_shape, dtype),
-          jtu.format_shape_dtype_string(out_shape, dtype)),
+          jtu.format_shape_dtype_string(out_shape, dtype),
+          order),
        "arg_shape": arg_shape, "out_shape": out_shape, "dtype": dtype,
-       "rng": jtu.rand_default()}
+       "order": order, "rng": jtu.rand_default()}
       for dtype in default_dtypes
+      for order in ["C", "F"]
       for arg_shape, out_shape in [
           (jtu.NUMPY_SCALAR_SHAPE, (1, 1, 1)),
           ((), (1, 1, 1)),
@@ -875,9 +877,9 @@ def fn(module, axis):
           ((2, 1, 4), (-1,)),
           ((2, 2, 4), (2, 8))
       ]))
-  def testReshape(self, arg_shape, out_shape, dtype, rng):
-    onp_fun = lambda x: onp.reshape(x, out_shape)
-    lnp_fun = lambda x: lnp.reshape(x, out_shape)
+  def testReshape(self, arg_shape, out_shape, dtype, order, rng):
+    onp_fun = lambda x: onp.reshape(x, out_shape, order=order)
+    lnp_fun = lambda x: lnp.reshape(x, out_shape, order=order)
     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)
@@ -1327,6 +1329,13 @@ def test_fail(x):
     result = api.grad(test_fail)(x)
     assert not onp.any(onp.isnan(result))
 
+  def testIssue453(self):
+    # https://github.com/google/jax/issues/453
+    a = onp.arange(6) + 1
+    ans = lnp.reshape(a, (3, 2), order='F')
+    expected = onp.reshape(a, (3, 2), order='F')
+    self.assertAllClose(ans, expected, check_dtypes=True)
+
 
 if __name__ == "__main__":
   absltest.main()