jax-ml/jax · error · TypeError

reshape new_sizes must all be positive, got {}.

Error message

reshape new_sizes must all be positive, got {}.

What it means

jax.lax.reshape requires every entry of new_sizes to be >= 0 (the error text says 'positive' but the check is d >= 0, allowing 0-sized dims); any negative dimension raises this TypeError. Unlike NumPy's reshape, lax.reshape does not accept a -1 wildcard — the target shape must be explicit.

Source

Thrown at jax/_src/lax/lax.py:7800

def shape_as_value(shape: core.Shape):
  """Converts a shape that may contain Poly values into a JAX value."""
  dtype = lax_utils.int_dtype_for_shape(shape, signed=True)
  if len(shape) == 0:
    return full((0,), np.array(0, dtype=dtype))
  if core.is_constant_shape(shape):
    return np.asarray(shape, dtype=dtype)
  dims = [
      expand_dims(convert_element_type(core.dimension_as_value(d), dtype),
                  (0,))
      for d in shape
  ]
  return concatenate(dims, dimension=0)

def _reshape_shape_rule(operand, *, new_sizes, dimensions, sharding):
  if not all(d >= 0 for d in new_sizes):
    msg = 'reshape new_sizes must all be positive, got {}.'
    raise TypeError(msg.format(new_sizes))
  # TODO(necula): re-enable this check
  if dimensions is not None:
    if set(dimensions) != set(range(np.ndim(operand))):
      msg = ('reshape dimensions must be a permutation of operand dimensions, '
             'got dimensions {} for shape {}.')
      raise TypeError(msg.format(dimensions, np.shape(operand)))
  return tuple(new_sizes)

class ReshapeExplicitError(Exception):
  pass

def _split_on_one_axis(op_shape, new_sizes):
  op_shape = [s for s in op_shape if s != 1]
  new_sizes = [s for s in new_sizes if s != 1]

  if len(new_sizes) <= len(op_shape):
    return False, []

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use jnp.reshape(x, (-1, 128)) or x.reshape(-1, 128), which supports -1 inference
2. Compute the flat size and derive dims explicitly: flat = x.size; rows = flat // 128
3. Validate all target dims are >= 0 before calling lax.reshape

Example fix

# before
y = jax.lax.reshape(x, new_sizes=[-1, 128])
# after
y = x.reshape(-1, 128)  # or jax.lax.reshape(x, (x.size // 128, 128))

assert all(d >= 0 for d in new_sizes), new_sizes
# for -1 inference, use jnp-level reshape or compute explicitly:
new_sizes = (x.size // known_dim, known_dim)

def valid_new_sizes(sizes) -> bool:
    return all(d >= 0 for d in sizes)

When it happens

Trigger: jax.lax.reshape(x, new_sizes=[-1, 128]) — using NumPy's infer-dimension idiom; computed target dims going negative.

Common situations: Porting numpy.reshape code with -1 directly to jax.lax.reshape; size arithmetic (e.g. n - k) producing negatives when k > n; dynamic batch inference code.

Related errors

• Dimension size after padding is not at least 0, got result s
• scan got `length` argument of {} which disagrees with leadin
• conv_general_dilated batch_group_count must divide lhs batch
• conv_general_dilated rhs output feature dimension size must
• conv_general_dilated window and window_strides must have the

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