jax-ml/jax · error · ValueError

dtype argument to `maxwell` must be a float dtype, got {dtyp

Error message

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

What it means

jax.random.maxwell (Maxwell-Boltzmann speed distribution) requires a floating-point dtype; it generates three standard normal components and takes the norm sqrt(X^2+Y^2+Z^2), which is float-only math. Integer or complex dtypes raise ValueError.

Source

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

      across devices in multi-device computation. Can be a
      :class:`~jax.sharding.NamedSharding`, a :class:`~jax.sharding.PartitionSpec`
      (``P``), or ``None`` (default). When specified, the output will be sharded
      according to the given sharding specification. Primarily used in explicit
      sharding mode.
      See the `explicit sharding tutorial <https://docs.jax.dev/en/latest/parallel.html>`_
      for more details.

  Returns:
    A jnp.array of samples, of shape `shape`.

  """
  # Generate samples using:
  # sqrt(X^2 + Y^2 + Z^2), X,Y,Z ~N(0,1)
  key, _ = _check_prng_key("maxwell", key)
  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 `maxwell` must be a float "
                     f"dtype, got {dtype}")
  shape = core.canonicalize_shape(shape)
  out_sharding = canonicalize_sharding_for_samplers(out_sharding, "maxwell", shape)
  return _maxwell(key, shape, dtype, out_sharding)


@jit(static_argnums=(1, 2, 3))
def _maxwell(key, shape, dtype, out_sharding) -> Array:
  shape = shape + (3,)
  if out_sharding is not None:
    new_partitions = (*out_sharding.spec, None)
    out_sharding = out_sharding.update(
        spec=out_sharding.spec.update(partitions=new_partitions))
  norm_rvs = normal(key=key, shape=shape, dtype=dtype, out_sharding=out_sharding)
  return jnp_linalg.norm(norm_rvs, axis=-1)


def double_sided_maxwell(key: ArrayLike,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pass jnp.float32/jnp.float64 or omit dtype (defaults to float).
2. If integer speeds are truly needed, sample float then cast afterwards: jax.random.maxwell(key, shape).astype(jnp.int32).
3. Validate configurable dtypes against np.floating before calling.

Example fix

// before
v = jax.random.maxwell(key, (1000, 3), dtype=jnp.int32)

// after
v = jax.random.maxwell(key, (1000, 3), dtype=jnp.float32)

from jax._src import dtypes
assert dtypes.issubdtype(dtypes.check_and_canonicalize_user_dtype(dtype or float), np.floating)

def is_float_dtype(dtype) -> bool:
    from jax._src import dtypes
    import numpy as np
    return dtypes.issubdtype(dtypes.check_and_canonicalize_user_dtype(dtype or float), np.floating)

When it happens

Trigger: jax.random.maxwell(key, shape, dtype=jnp.int32) or any dtype where dtypes.issubdtype(dtype, np.floating) is False.

Common situations: Molecular-dvelocity initialization scripts with a shared dtype constant; assuming a default int dtype; debugging configs that hard-code int32 globally.

Related errors

• dtype argument to `exponential` must be a float dtype, got {
• dtype argument to `gamma` must be a float dtype, got {dtype}
• dtype argument to `gumbel` must be a float dtype, got {dtype
• dtype argument to `laplace` must be a float dtype, got {dtyp
• dtype argument to `logistic` must be a float dtype, got {dty

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