jax-ml/jax · error · TypeError

Unknown FFT type value '{fft_type}'

Error message

Unknown FFT type value '{fft_type}'

What it means

The fft_type argument must be a str (parsed by name) or a lax.FftType enum member. Passing any other type (int, numpy integer, None) raises this TypeError, since JAX will not guess an enum from raw integers.

Source

Thrown at jax/_src/lax/fft.py:74

  if s in ("fft", "FFT"):
    return FftType.FFT
  elif s in ("ifft", "IFFT"):
    return FftType.IFFT
  elif s in ("rfft", "RFFT"):
    return FftType.RFFT
  elif s in ("irfft", "IRFFT"):
    return FftType.IRFFT
  else:
    raise ValueError(f"Unknown FFT type '{s}'")

@jit(static_argnums=(1, 2))
def fft(x, fft_type: FftType | str, fft_lengths: Sequence[int]):
  if isinstance(fft_type, str):
    typ = _str_to_fft_type(fft_type)
  elif isinstance(fft_type, FftType):
    typ = fft_type
  else:
    raise TypeError(f"Unknown FFT type value '{fft_type}'")

  if typ == FftType.RFFT:
    if np.iscomplexobj(x):
      raise ValueError("only real valued inputs supported for rfft")
    x = lax.convert_element_type(x, dtypes.to_inexact_dtype(dtypes.dtype(x)))
  else:
    x = lax.convert_element_type(x, dtypes.to_complex_dtype(dtypes.dtype(x)))
  if len(fft_lengths) == 0:
    # XLA FFT doesn't support 0-rank.
    return x
  fft_lengths = tuple(fft_lengths)
  return fft_p.bind(x, fft_type=typ, fft_lengths=fft_lengths)

def _fft_impl(x, fft_type, fft_lengths):
  return dispatch.apply_primitive(fft_p, x, fft_type=fft_type, fft_lengths=fft_lengths)

_complex_dtype = lambda dtype: (np.zeros((), dtype) + np.zeros((), np.complex64)).dtype
_real_dtype = lambda dtype: np.finfo(dtype).dtype

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Convert integers to lax.FftType(n) before calling, e.g. lax.FftType(0)
2. Store/transport the string names ('fft','ifft','rfft','irfft') instead of ints
3. Use jax.numpy.fft high-level functions which need no type argument

Example fix

# before
y = lax.fft(x, 0, (64,))
# after
y = lax.fft(x, lax.FftType(0), (64,))  # FFT type

if isinstance(fft_type, int):
    fft_type = lax.FftType(fft_type)
assert isinstance(fft_type, (str, lax.FftType))

def coerce_fft_type(t):
    if isinstance(t, int):
        return lax.FftType(t)
    if isinstance(t, str):
        return t
    raise TypeError(f'bad fft_type {t!r}')

When it happens

Trigger: lax.fft(x, 0, (n,)) using the XLA integer enum value, or passing an int pulled from a config file/serialized spec.

Common situations: Deserializing FFT configs where the type round-trips to a plain int; interoperating with code that uses XLA's numeric FftType.

Understand the failure class

Background: Invalid argument type errors: "must be of type string", "expected X, got Y", and ERR_INVALID_ARG_TYPE explained — this error's family across 15 libraries.

Related errors

• {name} was requested to map a value of non-array type {core.
• primal and tangent arguments to jax.jvp must be tuples or li
• {prim_name} takes a scalar pred as argument, got {pred}
• `compute_on`'s compute_type argument must be a string.
• Cannot interpret value of type {typ} as an abstract array; i

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