jax-ml/jax · error · ValueError

dtype argument to `binomial` must be a float dtype, got {dty

Error message

dtype argument to `binomial` must be a float dtype, got {dtype}

What it means

jax.random.binomial requires its dtype argument to be a floating-point dtype (default float), validated with dtypes.issubdtype(dtype, np.floating) after check_arraylike on n and p. Integer or bool dtypes raise this ValueError even though binomial counts are conceptually integral — JAX computes them in floating point.

Source

Thrown at jax/_src/random/core.py:3671

      representing the number of trials.
    p: a float or array of floats broadcast-compatible with ``shape``
      representing the probability of success of an individual trial.
    shape: optional, a tuple of nonnegative integers specifying the result
      shape. Must be broadcast-compatible with ``n`` and ``p``.
      The default (None) produces a result shape equal to ``np.broadcast(n, p).shape``.
    dtype: optional, a float dtype for the returned values (default float64 if
      jax_enable_x64 is true, otherwise float32).

  Returns:
    A random array with the specified dtype and with shape given by
    ``np.broadcast(n, p).shape``.
  """
  key, _ = _check_prng_key("binomial", key)
  check_arraylike("binomial", n, p)
  dtype = dtypes.check_and_canonicalize_user_dtype(
      float if dtype is None else dtype)
  if not dtypes.issubdtype(dtype, np.floating):
    raise ValueError(
        f"dtype argument to `binomial` must be a float dtype, got {dtype}"
      )
  if shape is not None:
    shape = core.canonicalize_shape(shape)
  return _binomial(key, n, p, shape, dtype)


# Functions related to key reuse checking
random_clone_p = core.Primitive("random_clone")
dispatch.simple_impl(random_clone_p)
random_clone_p.def_abstract_eval(lambda x: x)
batching.defvectorized(random_clone_p)
mlir.register_lowering(random_clone_p, lambda _, k: [k])


def multinomial(
    key: Array,
    n: RealArray,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Omit dtype or pass jnp.float32/np.float64
  2. Cast the result to int afterwards if needed: x.astype(jnp.int32)

Example fix

// before
b = jax.random.binomial(key, 10, 0.5, dtype=jnp.int32)
// after
b = jax.random.binomial(key, 10, 0.5, dtype=jnp.float32).astype(jnp.int32)
Defensive patterns

Strategy: validation

Validate before calling

import numpy as np
assert dtype is None or np.issubdtype(np.dtype(dtype).type, np.floating), 'binomial needs float dtype (cast result to int afterwards)'

Type guard

def is_float_dtype(d) -> bool:
    import numpy as np
    return d is None or np.issubdtype(np.dtype(d).type, np.floating)

Prevention

When it happens

Trigger: Calling jax.random.binomial(key, n, p, shape, dtype=np.int32) or dtype=np.bool_ despite n and p passing arraylike checks.

Common situations: Users assume a count-valued distribution should have an int dtype; porting NumPy/SciPy code that used int64 outputs.

Related errors


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