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} because the product of specified axis sizes ({math.prod(other_sizes)}) does not evenly divide {arr.size}

What it means

When reshaping with a single -1, JAX checks that the product of the explicitly-given axis sizes evenly divides the array's total size; otherwise the inferred dimension would be non-integral. This divisibility check only runs when all relevant sizes are concrete ints (not tracers).

Source

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

  """Fixes a -1 value in newshape, if present."""
  orig_newshape = newshape  # for error messages
  try:
    iter(newshape)  # pyrefly: ignore[no-matching-overload]
  except TypeError:
    newshape = [newshape]
  else:
    newshape: Sequence[DimSize]  # pyrefly: ignore[redefinition]
  newshape = core.canonicalize_shape(newshape)
  neg1s = [i for i, d in enumerate(newshape) if type(d) is int and d == -1]
  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`.

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Fix the target shape so sizes divide arr.size (e.g. use (-1, arr.size % k == 0 and k) or correct k)
  2. Print/inspect arr.shape and arr.size before reshaping in data-dependent code
  3. Use arr.reshape(arr.size // k, k) after asserting divisibility

Example fix

# before
x = jnp.zeros(10)
x.reshape(-1, 3)  # 10 % 3 != 0

# after
x.reshape(-1, 5)  # 10 % 5 == 0
Defensive patterns

Strategy: validation

Validate before calling

def reshape_infer(arr, k):
    assert arr.size % k == 0, f'{arr.size} not divisible by {k}'
    return arr.reshape(-1, k)

Prevention

When it happens

Trigger: `arr.reshape(-1, k)` where arr.size is not divisible by k, e.g. a size-10 array reshaped to (-1, 3).

Common situations: Hard-coded feature dimensions after a data change (e.g. flattened dims no longer match); off-by-one in sequence/window lengths feeding a reshape.

Related errors


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