jax-ml/jax · error · ValueError

dtype argument to `lognormal` must be a float or complex dty

Error message

dtype argument to `lognormal` must be a float or complex dtype, got {dtype}

What it means

jax.random.lognormal accepts any inexact dtype (float or complex), checked via dtypes.issubdtype(dtype, np.inexact). Integer and boolean dtypes fail this check and raise a ValueError; complex dtypes are allowed here unlike most samplers.

Source

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

      shape. The default (None) produces a result shape equal to ``()``.
    dtype: optional, a float dtype for the returned values (default float64 if
      jax_enable_x64 is true, otherwise float32).
    out_sharding: optional, specifies how the output array should be sharded
      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 random array with the specified dtype and with shape given by ``shape``.
  """
  key, _ = _check_prng_key("lognormal", key)
  dtype = dtypes.check_and_canonicalize_user_dtype(float if dtype is None else dtype)
  if not dtypes.issubdtype(dtype, np.inexact):
    raise ValueError(f"dtype argument to `lognormal` must be a float or complex dtype, "
                    f"got {dtype}")
  shape = _check_broadcast_shapes("lognormal", shape, sigma)
  out_sharding = canonicalize_sharding(out_sharding, "lognormal")
  _check_all_safe_to_cast("lognormal", dtype, sigma)
  return maybe_auto_axes(_lognormal, out_sharding, shape=shape, dtype=dtype)(key, sigma)

@jit(static_argnums=(2, 3), inline=True)
def _lognormal(key, sigma, shape, dtype) -> Array:
  sigma = lax.convert_element_type(sigma, dtype)
  scaled_norm = normal(key, shape, dtype) * sigma
  return lax.exp(scaled_norm)


def _stirling_approx_tail(k):
  stirling_tail_vals = jnp.array(
      [
          0.0810614667953272,
          0.0413406959554092,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Pass a float dtype (or complex if intended), or omit dtype
  2. Sample as float and round/cast afterwards if integer-like output is needed

Example fix

// before
x = jax.random.lognormal(key, sigma, dtype=jnp.int32)
// after
x = jax.random.lognormal(key, sigma, dtype=jnp.float32)
Defensive patterns

Strategy: validation

Validate before calling

import numpy as np
assert dtype is None or np.issubdtype(np.dtype(dtype).type, np.inexact), 'lognormal needs float or complex dtype'

Type guard

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

Prevention

When it happens

Trigger: Calling jax.random.lognormal(key, sigma, shape, dtype=np.int32) or dtype=np.bool_.

Common situations: Trying to generate lognormal counts as integers directly; complex output is supported but int is not, which surprises users porting NumPy code.

Related errors


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