jax-ml/jax · error · ValueError

corrcoef: dtype must be a subclass of float or complex; got

Error message

corrcoef: dtype must be a subclass of float or complex; got {dtype=}

What it means

jnp.corrcoef accepts a dtype parameter that must be an inexact (float or complex) dtype, since it is forwarded to cov for floating-point computation. A non-inexact dtype raises ValueError('corrcoef: dtype must be a subclass of float or complex; got {dtype=}') before any computation.

Source

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

           [-1.,  1.]], dtype=float32)

    The entries of the correlation matrix are normalized such that they
    lie within the range -1 to +1, where +1 indicates perfect correlation
    and -1 indicates perfect anti-correlation. For example, here is the
    correlation of 100 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.corrcoef(x))
    [[1.   0.03 0.12]
     [0.03 1.   0.01]
     [0.12 0.01 1.  ]]
  """
  util.check_arraylike("corrcoef", x)
  if dtype is not None and not dtypes.issubdtype(dtype, np.inexact):
    raise ValueError(f"corrcoef: dtype must be a subclass of float or complex; got {dtype=}")
  c = cov(x, y, rowvar, dtype=dtype)
  if len(np.shape(c)) == 0:
    # scalar - this should yield nan for values (nan/nan, inf/inf, 0/0), 1 otherwise
    return ufuncs.divide(c, c)
  d = diag(c)
  stddev = ufuncs.sqrt(ufuncs.real(d)).astype(c.dtype)
  c = c / stddev[:, None] / stddev[None, :]

  real_part = clip(ufuncs.real(c), -1, 1)
  if iscomplexobj(c):
    complex_part = clip(ufuncs.imag(c), -1, 1)
    c = lax.complex(real_part, complex_part)
  else:
    c = real_part
  return c


@export

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Pass a float dtype: dtype=jnp.float32 / jnp.float64
  2. Omit dtype entirely
  3. Guard with jnp.issubdtype(dtype, jnp.inexact) check

Example fix

// before
jnp.corrcoef(x, dtype=jnp.int32)
// after
jnp.corrcoef(x, dtype=jnp.float32)
Defensive patterns

Strategy: type-guard

Validate before calling

if dtype is not None:
    assert jnp.issubdtype(dtype, jnp.inexact), 'corrcoef dtype must be float/complex'
jnp.corrcoef(x, dtype=dtype)

Type guard

def is_inexact_dtype(d) -> bool:
    return d is None or jnp.issubdtype(d, jnp.inexact)

Prevention

When it happens

Trigger: jnp.corrcoef(x, dtype=jnp.int32) or any integer/bool dtype; passing np.int64 from a config.

Common situations: Users trying to control output precision assuming any dtype is allowed; templated code where dtype is injected and may be integral in some paths.

Related errors


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