jax-ml/jax · error · TypeError

series_order must be a Python integer.

Error message

series_order must be a Python integer.

What it means

jax.scipy.special.log_ndtr(x, series_order=3) requires series_order to be a plain Python int because it selects precomputed segment constants and controls unrolled series logic. Passing a float, np.float32, jax array, or other type raises TypeError.

Source

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

  `double-factorial
  <https://en.wikipedia.org/wiki/Double_factorial>`_ operator.


  Args:
    x: an array of type `float32`, `float64`.
    series_order: Positive Python integer. Maximum depth to
      evaluate the asymptotic expansion. This is the `N` above.

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

  Raises:
    TypeError: if `x.dtype` is not handled.
    TypeError: if `series_order` is a not Python `integer.`
    ValueError:  if `series_order` is not in `[0, 30]`.
  """
  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:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Coerce: log_ndtr(x, int(series_order))
2. Keep series_order as a static Python int; if using jit, mark it nondiff_argnums/static or bake it via functools.partial
3. Validate type early: assert isinstance(series_order, int)

Example fix

// before
jax.scipy.special.log_ndtr(x, order)  # order = 5.0 from config
// after
jax.scipy.special.log_ndtr(x, int(order))

series_order = int(series_order)
assert isinstance(series_order, int)

def py_int(v) -> bool:
    return isinstance(v, int) and not isinstance(v, bool)

When it happens

Trigger: Calling log_ndtr(x, 5.0), log_ndtr(x, np.int64(5)) on some versions, or passing a traced/computed value as series_order (e.g. from a config dict or hyperparameter sweep).

Common situations: Hyperparameter sweeps where order comes from argparse floats or JSON configs; passing a JAX scalar in place of a Python int; wrapping log_ndtr in vmap/pmap with order accidentally treated as an array.

Related errors

• series_order must be non-negative.
• series_order must be <= 30.
• x.dtype={np.dtype(dtype)} is not supported.
• Argument '{arg}' of type {type(arg)} is not a valid JAX type
• SymbolicScope constraint must be a string: got {repr(c_str)}

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