jax-ml/jax · error · ValueError

dtype argument to `double_sided_maxwell` must be a float dty

Error message

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

What it means

jax.random.double_sided_maxwell validates its optional dtype argument and only accepts floating-point dtypes. The dtype is canonicalized via dtypes.check_and_canonicalize_user_dtype (defaulting to Python float) and then checked with dtypes.issubdtype(dtype, np.floating). Passing any integer, bool, or complex dtype raises this ValueError before sampling.

Source

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

  where :math:`z = (x - \mu) / \sigma`, with the center :math:`\mu` specified by
  ``loc`` and the scale :math:`\sigma` specified by ``scale``.

  Args:
    key: a PRNG key.
    loc: The location parameter of the distribution.
    scale: The scale parameter of the distribution.
    shape: The shape added to the parameters loc and scale broadcastable shape.
    dtype: The type used for samples.

  Returns:
    A jnp.array of samples.

  """
  key, _ = _check_prng_key("double_sided_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 `double_sided_maxwell` must be a float"
                     f" dtype, got {dtype}")
  shape = core.canonicalize_shape(shape)
  return _double_sided_maxwell(key, loc, scale, shape, dtype)


@jit(static_argnums=(3, 4))
def _double_sided_maxwell(key, loc, scale, shape, dtype) -> Array:
  params_shapes = lax.broadcast_shapes(np.shape(loc), np.shape(scale))
  if not shape:
    shape = params_shapes

  shape = shape + params_shapes
  maxwell_key, rademacher_key = _split(key)
  maxwell_rvs = maxwell(maxwell_key, shape=shape, dtype=dtype)
  # Generate random signs for the symmetric variates.
  random_sign = rademacher(rademacher_key, shape=shape, dtype=dtype)
  assert random_sign.shape == maxwell_rvs.shape

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Pass a float dtype such as np.float32, jnp.float32, or np.float64, or omit dtype to use the default
  2. If you got the dtype from other data, cast it first: np.float32 if np.issubdtype(d, np.floating) else np.float32
  3. Check for accidental bool/int constants like dtype=0 or dtype=int

Example fix

// before
samples = jax.random.double_sided_maxwell(key, 0.0, 1.0, (1000,), dtype=np.int32)
// after
samples = jax.random.double_sided_maxwell(key, 0.0, 1.0, (1000,), dtype=np.float32)
Defensive patterns

Strategy: validation

Validate before calling

import numpy as np

def safe_double_sided_maxwell(key, loc, scale, shape, dtype=None):
    if dtype is not None and not np.issubdtype(np.dtype(dtype).type, np.floating):
        raise ValueError('dtype must be float; got %s' % dtype)
    return jax.random.double_sided_maxwell(key, loc, scale, shape, dtype)

Type guard

def is_float_dtype(dtype) -> bool:
    import numpy as np, jax.numpy as jnp
    if dtype is None or dtype is float: return True
    try: d = jnp.dtype(dtype)
    except TypeError: return False
    return np.issubdtype(d.type, np.floating)

Prevention

When it happens

Trigger: Calling jax.random.double_sided_maxwell(key, loc, scale, shape, dtype=np.int32), dtype=jnp.bfloat16 is fine but dtype=jnp.complex64 or any np.integer/np.bool_ dtype is not.

Common situations: Copy-pasting a dtype from another sampler (e.g. geometric which requires int), passing a dtype inferred from an integer array, or passing the string 'int32' instead of a float dtype.

Related errors


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