jax-ml/jax · error · TypeError

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

Error message

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

What it means

jax.scipy.special.spence (dilogarithm) accepts only float32 or float64 arrays. After jnp.asarray(x), lax.dtype(x) is checked against the whitelist and any other dtype (int, float16, bfloat16, complex) raises TypeError with the dtype interpolated.

Source

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

    function.

  Raises:
    TypeError: if elements of array `x` are not in (float32, float64).

  Notes:
    There is a different convention which defines Spence's function by the
    integral:

    .. math::

       -\int_0^x \frac{\log(1 - t)}{t}\mathrm{d}t

    This is our spence(1 - x).
  """
  x = jnp.asarray(x)
  dtype = lax.dtype(x)
  if dtype not in (np.float32, np.float64):
    raise TypeError(
      f"x.dtype={dtype} is not supported, see docstring for supported types.")
  return _spence(x)


def bernoulli(n: int) -> Array:
  r"""Generate the Bernoulli numbers :math:`B_0` through :math:`B_n`, inclusive.

  JAX implementation of :func:`scipy.special.bernoulli`.

  Args:
    n: integer, the index of the last Bernoulli number to generate.

  Returns:
    Array containing the Bernoulli numbers :math:`B_0` through :math:`B_n`, inclusive.

  Notes:
    ``bernoulli`` generates numbers using the :math:`B_n^-` convention,
    such that :math:`B_1=-1/2`.

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Cast explicitly: spence(jnp.asarray(x, dtype=jnp.float64)) (enable x64 first if needed).
2. Use float32 for TPU/GPU pipelines; avoid bfloat16 around special functions.
3. Wrap calls in a helper that canonicalizes dtype.

Example fix

# before (raises)
y = jax.scipy.special.spence(2)

# after
y = jax.scipy.special.spence(jnp.asarray(2, dtype=jnp.float64))

x = jnp.asarray(x)
if x.dtype not in (jnp.float32, jnp.float64):
    x = x.astype(jnp.float64 if jax.config.x64_enabled else jnp.float32)

def as_supported_float(x):
    x = jnp.asarray(x)
    return x if x.dtype in (jnp.float32, jnp.float64) else x.astype(jnp.float32)

When it happens

Trigger: Calling jax.scipy.special.spence with integer inputs (e.g. spence(2)), bfloat16/float16 tensors, or complex values.

Common situations: Convenience calls with Python ints; half-precision pipelines on TPU/GPU; forgetting that asarray does not auto-promote ints to float here.

Related errors

• x argument to bincount must have an integer type; got {x.dty
• primal and tangent arguments to jax.jvp do not match; dtypes
• unexpected JAX type (e.g. shape/dtype) for gradient ref pass
• Accumulator aval mismatch: expected {aval}, got {acc.aval}
• unexpected JAX type (e.g. shape/dtype) for argument to VJP f

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