jax-ml/jax · error · ValueError

jax.numpy.meshgrid only supports copy=True

Error message

jax.numpy.meshgrid only supports copy=True

What it means

jnp.meshgrid does not support copy=False. NumPy's meshgrid can return views when copy=False, but JAX arrays have different view semantics under transformations, so JAX only implements the copying behavior and raises this ValueError for any falsy copy argument.

Source

Thrown at jax/_src/numpy/lax_numpy.py:6092

    [[1 2]]
    >>> print(y_grid)
    [[10]
     [20]
     [30]]

    2D matrix-index mesh grid:

    >>> x_grid, y_grid = jnp.meshgrid(x, y, indexing='ij')
    >>> print(x_grid)
    [[1 1 1]
     [2 2 2]]
    >>> print(y_grid)
    [[10 20 30]
     [10 20 30]]
  """
  args = list(util.ensure_arraylike_tuple("meshgrid", tuple(xi)))
  if not copy:
    raise ValueError("jax.numpy.meshgrid only supports copy=True")
  if indexing not in ["xy", "ij"]:
    raise ValueError(f"Valid values for indexing are 'xy' and 'ij', got {indexing}")
  if any(a.ndim != 1 for a in args):
    raise ValueError("Arguments to jax.numpy.meshgrid must be 1D, got shapes "
                     f"{[a.shape for a in args]}")
  if indexing == "xy" and len(args) >= 2:
    args[0], args[1] = args[1], args[0]
  shape = [1 if sparse else a.shape[0] for a in args]
  _a_shape = lambda i, a: [*shape[:i], a.shape[0], *shape[i + 1:]] if sparse else shape
  output = [lax.broadcast_in_dim(a, _a_shape(i, a), (i,)) for i, a, in enumerate(args)]
  if indexing == "xy" and len(args) >= 2:
    output[0], output[1] = output[1], output[0]
  return tuple(output)


@export
@api.jit
def i0(x: ArrayLike) -> Array:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Drop copy=False (use the default copy=True) and add sparse=True if memory is the concern
2. If only coordinate values matter, use jnp.broadcast_to on 1-D arrays to get logical grids without materialization

Example fix

// before
X, Y = jnp.meshgrid(x, y, copy=False)
// after
X, Y = jnp.meshgrid(x, y, sparse=True)

out = jnp.meshgrid(*xi, sparse=True)  # never pass copy=False

When it happens

Trigger: Calling jnp.meshgrid(*xi, copy=False), typically ported from np.meshgrid(..., copy=False, sparse=True) memory-optimization patterns.

Common situations: Porting NumPy code that disabled copying to save memory on large grids; users assuming JAX supports the full np.meshgrid signature. Note jnp.meshgrid does support sparse=True for the same memory goal.

Related errors

• Invalid value "{default}" for JAX flag {name}
• new enum value must be in {enum_values}, got {new_val} of ty
• new enum value must be None or in {enum_values}, got {new_va
• Invalid value "{default_str}" for JAX flag {name}
• Invalid value "{default_env}" for JAX flag {name}

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