jax-ml/jax · error · TypeError

x.dtype={} is not supported, see docstring for supported typ

Error message

x.dtype={} is not supported, see docstring for supported types.

What it means

jax.scipy.special.ndtr (standard normal CDF) only accepts float32 or float64 arrays. After asarray, any other dtype (int, complex, bfloat16, float16) raises TypeError pointing at the docstring's supported types.

Source

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

    \mathrm{ndtr}(x) =&
      \ \frac{1}{\sqrt{2 \pi}}\int_{-\infty}^{x} e^{-\frac{1}{2}t^2} \mathrm{d}t \\
    =&\ \frac{1}{2} (1 + \mathrm{erf}(\frac{x}{\sqrt{2}})) \\
    =&\ \frac{1}{2} \mathrm{erfc}(-\frac{x}{\sqrt{2}})
    \end{align}

  Args:
    x: An array of type `float32`, `float64`.

  Returns:
    An array with `dtype=x.dtype`.

  Raises:
    TypeError: if `x` is not floating-type.
  """
  x = jnp.asarray(x)
  dtype = lax.dtype(x)
  if dtype not in (np.float32, np.float64):
    raise TypeError(
        "x.dtype={} is not supported, see docstring for supported types."
        .format(dtype))
  return _ndtr(x)


def ndtri(p: ArrayLike) -> Array:
  r"""The inverse of the CDF of the Normal distribution function.

  JAX implementation of :obj:`scipy.special.ndtri`.

  Returns `x` such that the area under the PDF from :math:`-\infty` to `x` is equal
  to `p`.

  A piece-wise rational approximation is done for the function.
  This is based on the implementation in netlib.

  Args:
    p: an array of type `float32`, `float64`.

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Cast to float: ndtr(x.astype(jnp.float32)) or jnp.asarray(x, jnp.float64)
Use promote_args_inexact-style casting upstream so tensors entering statistical functions are already float32/64
For bfloat16 pipelines, keep ndtr in float32 and cast back

Example fix

// before
jax.scipy.special.ndtr(jnp.array([0, 1, 2]))  # int32
// after
jax.scipy.special.ndtr(jnp.array([0, 1, 2], dtype=jnp.float32))

Defensive patterns

Strategy: validation

Validate before calling

x = jnp.asarray(x)
if x.dtype not in (jnp.float32, jnp.float64):
    x = x.astype(jnp.float32)
ndtr(x)

Type guard

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

Prevention

Always cast inputs to statistical functions to float32/64 at pipeline boundaries
Avoid bf16/fp16 tensors reaching special functions in mixed-precision training

When it happens

Trigger: Calling ndtr on integer arrays (ndtr(jnp.arange(3))), bfloat16/float16 tensors from TPUs or mixed-precision training, or complex inputs.

Common situations: Passing integer-valued data (counts, bins) into a Gaussian CDF; half-precision model weights under jit with x64 disabled; assuming automatic casting like NumPy's special functions.

Related errors

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