jax-ml/jax · error · ValueError

m has more than 2 dimensions

Error message

m has more than 2 dimensions

What it means

jnp.cov requires the primary data matrix m to be 1D (a single variable's observations) or 2D (variables x observations). If m.ndim > 2 after promotion, ValueError('m has more than 2 dimensions') is raised, matching NumPy's behavior.

Source

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

    points drawn from a 3-dimensional standard normal distribution:

    >>> key = jax.random.key(0)
    >>> x = jax.random.normal(key, shape=(3, 100))
    >>> with jnp.printoptions(precision=2):
    ...   print(jnp.cov(x))
    [[0.9  0.03 0.1 ]
     [0.03 1.   0.01]
     [0.1  0.01 0.85]]
  """
  if y is not None:
    m, y = util.promote_args_inexact("cov", m, y)
    if y.ndim > 2:
      raise ValueError("y has more than 2 dimensions")
  else:
    m, = util.promote_args_inexact("cov", m)

  if m.ndim > 2:
    raise ValueError("m has more than 2 dimensions")  # same as numpy error

  if dtype is not None and not dtypes.issubdtype(dtype, np.inexact):
    raise ValueError(f"cov: dtype must be a subclass of float or complex; got {dtype=}")

  X = atleast_2d(m)
  if not rowvar and m.ndim != 1:
    X = X.T
  if X.shape[0] == 0:
    return array([]).reshape(0, 0)

  if y is not None:
    y_arr = atleast_2d(y)
    if not rowvar and y_arr.shape[0] != 1:
      y_arr = y_arr.T
    X = concatenate((X, y_arr), axis=0)
  if X.shape[1] == 0:
    cov_shape = () if X.shape[0] == 1 else (X.shape[0], X.shape[0])
    return array_creation.full(cov_shape, np.nan, dtype=X.dtype)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Flatten leading dimensions: m.reshape(-1, m.shape[-1]) or m.reshape(m.shape[0], -1) depending on variable layout
  2. Use jax.vmap(jnp.cov) to get per-sample covariance matrices for batched data
  3. Restructure data so rows are variables and columns observations

Example fix

// before
c = jnp.cov(batch_3d)  # ValueError
// after
c = jax.vmap(jnp.cov)(batch_3d)  # or reshape to 2D
Defensive patterns

Strategy: validation

Validate before calling

m = jnp.asarray(m)
assert m.ndim <= 2, 'm must be 1D or 2D for cov'
jnp.cov(m)

Type guard

def is_cov_input(x) -> bool:
    return 1 <= jnp.asarray(x).ndim <= 2

Prevention

When it happens

Trigger: Calling jnp.cov on a 3D tensor such as shape (batch, features, time), e.g. jnp.cov(images) where images.ndim == 3.

Common situations: Applying cov to batches of images, video, or windowed time-series without flattening; migrating NumPy pipelines that already reshaped data but losing the reshape step.

Related errors


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