jax-ml/jax · error · TypeError

cannot reshape array of shape {arr.shape} (size {arr.size})

Error message

cannot reshape array of shape {arr.shape} (size {arr.size}) into shape {orig_newshape} (size {math.prod(newshape)})

What it means

For a fully-specified new shape, `_compute_newshape` verifies that the product of the new dimensions equals the array's total size (when all sizes are concrete ints). NumPy has the same rule ('cannot reshape array of size N into shape ...'); JAX phrases it with shape and size details.

Source

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

  if len(neg1s) > 1:
    raise TypeError("can only specify one unknown axis size with a `-1` value, "
                    f"got {orig_newshape}")
  if neg1s:
    i, = neg1s
    other_sizes = (*newshape[:i], *newshape[i+1:])
    if (all(isinstance(d, int) for d in (*arr.shape, *other_sizes)) and
        arr.size % math.prod(other_sizes) != 0):
      raise TypeError(f"cannot reshape array of shape {arr.shape} (size {arr.size}) "
                      f"into shape {orig_newshape} because the product of "
                      f"specified axis sizes ({math.prod(other_sizes)}) does "
                      f"not evenly divide {arr.size}")
    sz = core.cancel_divide_tracers(arr.shape, other_sizes)
    if sz is not None:
      return (*newshape[:i], sz, *newshape[i+1:])
  else:
    if (all(isinstance(d, int) for d in (*arr.shape, *newshape)) and
        arr.size != math.prod(newshape)):
      raise TypeError(f"cannot reshape array of shape {arr.shape} (size {arr.size}) "
                      f"into shape {orig_newshape} (size {math.prod(newshape)})")
  return tuple(-core.divide_shape_sizes(arr.shape, newshape)
               if core.definitely_equal(d, -1) else d for d in newshape)

def _view(self: Array, dtype: DTypeLike | None = None, type: None = None) -> Array:
  """Return a bitwise copy of the array, viewed as a new dtype.

  This is fuller-featured wrapper around :func:`jax.lax.bitcast_convert_type`.

  If the source and target dtype have the same bitwidth, the result has the same
  shape as the input array. If the bitwidth of the target dtype is different
  from the source, the size of the last axis of the result is adjusted
  accordingly.

  >>> jnp.zeros([1,2,3], dtype=jnp.int16).view(jnp.int8).shape
  (1, 2, 6)
  >>> jnp.zeros([1,2,4], dtype=jnp.int8).view(jnp.int16).shape
  (1, 2, 2)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Check arr.shape/arr.size and update the target shape
  2. Derive one dimension from the data: arr.reshape(arr.shape[0], -1)
  3. Add an assert on the incoming shape early in the pipeline to fail near the cause

Example fix

# before
imgs.reshape(32, 28, 28)  # imgs.shape == (32, 784)

# after
imgs.reshape(32, 28, 28) if imgs.ndim == 2 else imgs  # or
imgs.reshape(imgs.shape[0], 28, 28)
Defensive patterns

Strategy: validation

Validate before calling

def reshape_exact(arr, shape):
    import math
    assert arr.size == math.prod(s for s in shape if s != -1) or math.prod(shape) == arr.size
    return arr.reshape(shape)

Prevention

When it happens

Trigger: `arr.reshape(2, 3)` where arr.size != 6, e.g. reshaping a size-12 array into (2, 3).

Common situations: Data shape changes upstream (different batch size, image resolution, channels-last vs channels-first); stale hard-coded shapes in model input pipelines.

Related errors


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