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
- Drop the order argument (default 'C' matches JAX semantics)
- Use order='F' if Fortran order was intended
- 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
- Never propagate order='A' into JAX
- Strip/normalize the order kwarg in NumPy/JAX-agnostic wrappers
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
- Only implemented for order='K'
- `type` argument of array.view() is not supported.
- JAX Arrays do not implement the arr.flat property: consider
- array ref with memory space only works inside of a `jit`.
- pinned array ref only works inside of a `jit`.
AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27).
Data as JSON: /api/errors/b555d8c7b8cd1b3c.
Report an issue: GitHub.