jax-ml/jax · error · ValueError

x and n must be of integer type; got x.dtype={x.dtype}, n.dt

Error message

x and n must be of integer type; got x.dtype={x.dtype}, n.dtype={n.dtype}

What it means

jax.scipy.stats.multinomial.logpmf (and pmf which calls it) requires the count vector x and total count n to have integer dtypes. After promotion via promote_args_numeric, the code checks dtypes.issubdtype(x.dtype, np.integer) and rejects floating-point inputs. This mirrors scipy's requirement that multinomial counts be integers, since non-integer counts are mathematically undefined for a discrete distribution.

Source

Thrown at jax/_src/scipy/stats/multinomial.py:52

     f(x, n, p) = n! \prod_{i=1}^k \frac{p_i^{x_i}}{x_i!}

  with :math:`n = \sum_i x_i`.

  Args:
    x: arraylike, value at which to evaluate the PMF
    n: arraylike, distribution shape parameter
    p: arraylike, distribution shape parameter

  Returns:
    array of logpmf values.

  See Also:
    :func:`jax.scipy.stats.multinomial.pmf`
  """
  p, = promote_args_inexact("multinomial.logpmf", p)
  x, n = promote_args_numeric("multinomial.logpmf", x, n)
  if not dtypes.issubdtype(x.dtype, np.integer):
    raise ValueError(f"x and n must be of integer type; got x.dtype={x.dtype}, n.dtype={n.dtype}")
  x = x.astype(p.dtype)
  n = n.astype(p.dtype)
  logprobs = gammaln(n + 1) + jnp.sum(xlogy(x, p) - gammaln(x + 1), axis=-1)
  return jnp.where(jnp.equal(jnp.sum(x), n), logprobs, -np.inf)


def pmf(x: ArrayLike, n: ArrayLike, p: ArrayLike) -> Array:
  r"""Multinomial probability mass function.

  JAX implementation of :obj:`scipy.stats.multinomial` ``pmf``.

  The multinomial probability distribution is given by

  .. math::

     f(x, n, p) = n! \prod_{i=1}^k \frac{p_i^{x_i}}{x_i!}

  with :math:`n = \sum_i x_i`.

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Cast x and n to an integer dtype before calling: x.astype(jnp.int32), n=int(n) or jnp.asarray(n, dtype=jnp.int32)
2. Verify your count data actually represents integer counts; if x holds probabilities instead of counts, you're calling the wrong function
3. Ensure n equals sum(x) along the last axis, otherwise the result is -inf even with correct dtypes

Example fix

// before
p = jnp.array([0.5, 0.5])
x = jnp.array([1.0, 1.0])  # float -> raises
n = 2.0
jax.scipy.stats.multinomial.logpmf(x, n, p)

// after
x = jnp.array([1, 1])
n = 2
jax.scipy.stats.multinomial.logpmf(x, n, p)

import numpy as np
import jax.numpy as jnp

def check_multinomial_inputs(x, n):
    x, n = jnp.asarray(x), jnp.asarray(n)
    assert np.issubdtype(x.dtype, np.integer), f"x must be int, got {x.dtype}"
    assert np.issubdtype(n.dtype, np.integer), f"n must be int, got {n.dtype}"
    return x, n

def is_int_array(a) -> bool:
    return jnp.issubdtype(jnp.asarray(a).dtype, jnp.integer)

try:
    lp = multinomial.logpmf(x, n, p)
except ValueError as e:
    if 'integer type' in str(e):
        x, n = x.astype(jnp.int32), int(n)
        lp = multinomial.logpmf(x, n, p)
    else: raise

When it happens

Trigger: Calling jax.scipy.stats.multinomial.logpmf or .pmf with x or n as float arrays, e.g. x=jnp.array([1.0, 2.0]) or n=10.0, or passing Python floats that promote to float32/float64. Also occurs when data loaded from float sources (e.g. CSVs, normalized probabilities) is passed as counts.

Common situations: Users coming from continuous distributions, data pipelines that produce float arrays by default, or JAX's x64-disabled mode where integer division produces floats. Version changes that made dtype checking stricter also surface latent float inputs.

Related errors

• dtype must be a complex floating-point type; got {dtype}.
• primal and tangent arguments to jax.jvp do not match; dtypes
• unexpected JAX type (e.g. shape/dtype) for gradient ref pass
• Accumulator aval mismatch: expected {aval}, got {acc.aval}
• unexpected JAX type (e.g. shape/dtype) for argument to VJP f

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