jax-ml/jax · error · TypeError

x.dtype={np.dtype(dtype)} is not supported.

Error message

x.dtype={np.dtype(dtype)} is not supported.

What it means

log_ndtr only supports float32 and float64 inputs; the implementation selects precomputed lower/upper segment tables keyed on those dtypes. Any other dtype (int, complex, bfloat16, float16) raises TypeError with the offending dtype.

Source

Thrown at jax/_src/scipy/special.py:1682

  """
  if not isinstance(series_order, int):
    raise TypeError("series_order must be a Python integer.")
  if series_order < 0:
    raise ValueError("series_order must be non-negative.")
  if series_order > 30:
    raise ValueError("series_order must be <= 30.")

  x_arr = jnp.asarray(x)
  dtype = lax.dtype(x_arr)

  if dtype == np.float64:
    lower_segment: np.ndarray = _LOGNDTR_FLOAT64_LOWER
    upper_segment: np.ndarray = _LOGNDTR_FLOAT64_UPPER
  elif dtype == np.float32:
    lower_segment = _LOGNDTR_FLOAT32_LOWER
    upper_segment = _LOGNDTR_FLOAT32_UPPER
  else:
    raise TypeError(f"x.dtype={np.dtype(dtype)} is not supported.")

  # The basic idea here was ported from:
  #   https://root.cern.ch/doc/v608/SpecFuncCephesInv_8cxx_source.html
  # We copy the main idea, with a few changes
  # * For x >> 1, and X ~ Normal(0, 1),
  #     Log[P[X < x]] = Log[1 - P[X < -x]] approx -P[X < -x],
  #     which extends the range of validity of this function.
  # * We use one fixed series_order for all of 'x', rather than adaptive.
  # * Our docstring properly reflects that this is an asymptotic series, not a
  #   Taylor series. We also provided a correct bound on the remainder.
  # * We need to use the max/min in the _log_ndtr_lower arg to avoid nan when
  #   x=0. This happens even though the branch is unchosen because when x=0
  #   the gradient of a select involves the calculation 1*dy+0*(-inf)=nan
  #   regardless of whether dy is finite. Note that the minimum is a NOP if
  #   the branch is chosen.
  x_arr_gt_upper_segment = lax.gt(x_arr, upper_segment)
  ndtr_arg = jnp.where(x_arr_gt_upper_segment, -x_arr,
                       lax.max(x_arr, lower_segment))

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Cast: log_ndtr(jnp.asarray(x, jnp.float32))
2. Promote to inexact dtype upstream so statistics always see float32/64
3. In bf16 pipelines, upcast around log_ndtr and downcast after

Example fix

// before
jax.scipy.special.log_ndtr(x)  # x is bfloat16
// after
jax.scipy.special.log_ndtr(x.astype(jnp.float32))

x = jnp.asarray(x, jnp.float32)  # or float64
log_ndtr(x, series_order)

def is_f32_f64(x):
    return jnp.dtype(x) in (np.float32, np.float64)

When it happens

Trigger: Calling log_ndtr(jnp.array([0, 1])) (int), log_ndtr on bfloat16/float16 tensors, or complex x. Reached via logcdf and _log_gauss_mass in distribution code with non-float inputs.

Common situations: log-probability computations on integer-encoded data; mixed-precision (bf16) training calling logcdf of a normal; assuming complex support for log Φ of complex arguments.

Related errors

• betainc does not support complex-valued inputs.
• Argument `n` to polygamma must be of integer type. Got dtype
• x.dtype={} is not supported, see docstring for supported typ
• series_order must be a Python integer.
• series_order must be non-negative.

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