Dependencies

Latest stable releases of major dependencies (as of 2026-10-13):

Library	Version
JAX	0.11.2
Flax	0.12.10
Equinox	0.13.8
onnx-ir	1.0.0
onnx	1.23.2
onnxruntime	1.31.0
onnxruntime-web	1.30.0

The latest releases are not necessarily selected for every supported Python version; Poetry resolves compatible versions for each environment. The supported ONNX IR minimum is 1.0.0 (onnx-ir>=1.0.0). onnxruntime-web 1.30.0 is locked in package-lock.json for the Node.js and Chromium WASM validation flows. The supported Python runtime minimum remains ONNX Runtime 1.24.1; exports do not select operators based on the runtime installed on the export machine.

ONNX Runtime kernel gaps

Some lowerings depend on kernels or behavior that ONNX Runtime lacks. Because exports must run on every supported version, a workaround can be removed only once the supported minimum includes the fix; a fix in the locked release alone is not enough. Re-check this table whenever the locked ONNX Runtime or the supported minimum changes, and update the versions it was checked against. The minimum-dependencies CI job runs the related precision tests at the supported minimum.

Gap	Checked	Affected exports	Current handling
No bfloat16 Log on CPU or CUDA; no bfloat16 Exp on CPU	CPU 1.24.1, 1.31.0; CUDA 1.31.0	bfloat16 log1p, expm1, sinh, jax.nn.log1mexp, jax.nn.log_sigmoid	Computed in float32 and rounded back. Float32 is also the more accurate choice (within 1 ulp), so a new kernel alone is no reason to switch.
No float64 or bfloat16 Sinh on CPU; no Sinh on CUDA	CPU 1.24.1, 1.31.0; CUDA 1.31.0	jnp.sinh; lax.sinh other than float32	expm1-based form. Float32 lax.sinh keeps Sinh, which runs on CPU in CUDA sessions.
No float64 Cos on CPU	Missing in 1.24.1, present in 1.31.0	float64 lax.cos, jnp.cos	Half-angle form 1 - 2 * Sin(x / 2)²
No float64 Erf on CPU	CPU 1.24.1, 1.31.0	float64 exact jax.nn.gelu, nnx.gelu	None: the export uses Erf (or Gelu, which expands to it) and fails to load; float64 test variants are skipped.
CPU Where returns +0.0 when it selects -0.0	CPU 1.24.1, 1.31.0	log1p, expm1, sinh	log1p and sinh end in Mul to keep the sign; expm1(-0.0) exports as +0.0.
MatMulInteger with a zero contraction length can leave output values uninitialized	CPU 1.31.0	INT8/UINT8 lax.dot_general with INT32 accumulation	A known empty contraction produces explicit INT32 zeros. Potentially empty symbolic contractions pad both operands with one zero product; this may copy the operands. Fixed positive contractions retain direct integer multiplication.
UINT8 × INT8 MatMulInteger can saturate intermediate pair sums on some x86 CPUs	CPU 1.31.0 on GitHub runners	UINT8-left/INT8-right lax.dot_general with INT32 accumulation	Re-encode the right operand as UINT8 by adding 128 in INT32, then use a UINT8 right zero point of 128. This keeps the full source range and exact integer accumulation.
BF16 Clip and Where kernels are unavailable	CPU 1.31.0	BF16 conversion to INT8/UINT8	Widen exactly to FP32 for NaN handling and saturation, then cast to the requested integer type. FP16, FP32 and FP64 retain their actual computation type. Minimum-runtime qualification is separate.

ONNX Runtime documents the UINT8 × INT8 saturation mechanism and UINT8 × UINT8 workaround. The affected AVX2/AVX512 kernels use saturating INT16 intermediate sums; VNNI and Arm paths do not have that issue. A passing local runtime test alone therefore does not establish correctness on the other CPU paths.

An ONNX Cast alone also differs from the observed JAX FP16/BF16/FP32/FP64 conversion to INT8/UINT8 at overflow, infinity and NaN boundaries. Those conversions map NaN to zero and clamp to the destination range before casting. Integer-to-integer conversions keep their existing wrapping behaviour.

To check a kernel, create a session for a one-node model with that operator and dtype. For providers other than CPU, set the session option session.disable_cpu_ep_fallback to 1 so the node cannot fall back to CPU. A NOT_IMPLEMENTED error means the kernel is missing.

For dependency requirements and optional dependency groups, see pyproject.toml. For resolved Python dependency versions, see poetry.lock.