jax-ml/jax · error · NotImplementedError

np.reshape order=A is not implemented.

Error message

np.reshape order=A is not implemented.

What it means

`jax.numpy.reshape` (and `Array.reshape`) supports order='C' and order='F' but not order='A' (Fortran-contiguous-if-input-is-F-else-C), because JAX arrays have no contiguity concept. NumPy accepts 'A'; JAX raises NotImplementedError.

Source

Thrown at jax/_src/numpy/array_methods.py:386

def _reshape(self: Array, *args: Any, order: str = "C", out_sharding=None
             ) -> Array:
  """Returns an array containing the same data with a new shape.

  Refer to :func:`jax.numpy.reshape` for full documentation.
  """
  __tracebackhide__ = True
  newshape = _compute_newshape(self, args[0] if len(args) == 1 else args)
  if order == "C":
    return lax.reshape(self, newshape, None, out_sharding=out_sharding)
  elif order == "F":
    dims = list(range(self.ndim)[::-1])
    out_sharding = canonicalize_sharding(out_sharding, "jnp.reshape")
    out_sharding = (
        None if out_sharding is None else out_sharding.update(
            spec=out_sharding.spec.update(partitions=out_sharding.spec[::-1])))
    return lax.reshape(self, newshape[::-1], dims, out_sharding=out_sharding).T
  elif order == "A":
    raise NotImplementedError("np.reshape order=A is not implemented.")
  else:
    raise ValueError(f"Unexpected value for 'order' argument: {order}.")

def _round(self: Array, decimals: int = 0, out: None = None) -> Array:
  """Round array elements to a given decimal.

  Refer to :func:`jax.numpy.round` for full documentation.
  """
  return lax_numpy.round(self, decimals=decimals, out=out)

def _searchsorted(self: Array, v: ArrayLike, side: str = 'left',
                  sorter: ArrayLike | None = None, *, method: str = 'scan') -> Array:
  """Perform a binary search within a sorted array.

  Refer to :func:`jax.numpy.searchsorted` for full documentation."""
  return lax_numpy.searchsorted(self, v, side=side, sorter=sorter, method=method)

def _sort(self: Array, axis: int | None = -1, *, kind: None = None,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Drop the order argument (default 'C' matches JAX semantics)
  2. Use order='F' if Fortran order was intended
  3. Branch on order != 'A' before calling reshape in shared NumPy/JAX code

Example fix

# before
jnp.reshape(a, (2, 3), order='A')

# after
jnp.reshape(a, (2, 3))  # or order='F' if that was meant
Defensive patterns

Strategy: validation

Validate before calling

def reshape_compat(a, shape, order='C'):
    if order == 'A':
        order = 'C'
    return a.reshape(shape, order=order)

Prevention

When it happens

Trigger: Calling `arr.reshape(shape, order='A')` or `jnp.reshape(arr, shape, order='A')`.

Common situations: Porting NumPy code that passes order='A'; generic library code where order is a parameter that can be 'A'.

Related errors


AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/b555d8c7b8cd1b3c. Report an issue: GitHub.