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.

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 c29397b0b..97f0f7f72 100644
--- 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"):
diff --git a/tests/lax_numpy_test.py b/tests/lax_numpy_test.py
index f4cb9753c..8bc59c6dd 100644
--- a/tests/lax_numpy_test.py
+++ b/tests/lax_numpy_test.py
@@ -859,12 +859,14 @@ class LaxBackedNumpyTests(jtu.JaxTestCase):
     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 @@ class LaxBackedNumpyTests(jtu.JaxTestCase):
           ((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 @@ class LaxBackedNumpyTests(jtu.JaxTestCase):
     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()