[dependabot]

Updates the requirements on jax2onnx to permit the latest version.

Release notes

Sourced from jax2onnx's releases.

0.17.0

jax2onnx 0.17.0 — accuracy-first normalization, faithful operator exports, modern JAX compatibility

• Changed the default normalization policy: Switched the specialized Equinox and Flax NNX/Linen LayerNorm exports, together with Flax NNX/Linen RMSNorm, to explicit graphs under normalization_mode="auto". This extends the existing explicit defaults for Flax NNX/Linen GroupNorm and Equinox RMSNorm. Export applies a preselected plugin policy without benchmarking individual models. The documented qualification policy requires any native replacement for auto to meet the same fixed acceptance bounds. Fixed comparative bounds currently cover selected float32 LayerNorm cases on ONNX Runtime CPU; equivalent GroupNorm and RMSNorm bounds remain future work. "prefer_native" requests an eligible native operator without requiring it to meet auto’s accuracy bounds, while "force_decomposed" requests the explicit graph.

• Made LayerNorm exports faithful to framework statistics: Added normalization_mode support to the specialized Equinox and Flax NNX/Linen LayerNorm exporters, preserving two-pass variance with a constant-row safeguard for Equinox and slow-variance NNX, or Flax’s clamped fast variance. Linen slow-variance LayerNorm continues to trace JAX in every mode. The tested explicit paths remain explicit under ONNX Runtime CPU optimizations and are checked against stricter outlier-input accuracy bounds than the native path. Exports below opset 17 no longer emit LayerNormalization. Hardened the constant-row and constant-group safeguards in explicit slow-variance LayerNorm and GroupNorm so nonfinite inputs retain JAX’s nonfinite results.

• Expanded RMSNorm export options and protected explicit paths: Added normalization_mode support for eqx.nn.RMSNorm, matching Flax NNX/Linen RMSNorm: "prefer_native" can emit ONNX RMSNormalization at opset 23 or newer, with an additional Add for Equinox’s optional bias. Other modes and older opsets retain the explicit graph. Equinox and Flax NNX/Linen explicit RMSNorm graphs now square with Mul instead of Pow. Targeted ONNX Runtime CPU tests confirm that the checked Equinox and NNX paths avoid fusion into native normalization kernels.

• Preserved transposed-convolution geometry: Corrected ONNX ConvTranspose lowering for input-dilated lax.conv_general_dilated, covering eqx.nn.ConvTranspose, nnx.ConvTranspose, jax.lax.conv_transpose, and nnx.Conv with input_dilation. The lowering maps input dilation to strides, uses a group-aware kernel layout, and inserts an explicit output Pad where JAX pads beyond the kernel’s reach, preserving JAX shapes and values. Input-dilated convolutions now fail explicitly when combined with non-unit window_strides or complex values; all lax.conv_general_dilated lowerings reject batch_group_count > 1.

• Made GELU exports opset-aware: Restricted native ONNX Gelu to opset 20 and newer. Below opset 20, jax.nn.gelu and nnx.gelu now use the exact Erf formula or the tanh approximation, avoiding an unsupported Gelu node at older opsets.

• Maintained JIT and gradient compatibility with newer JAX releases: Updated JIT lowering to create fresh variables through the compatibility layer, supporting the JAX 0.11.1+ Var constructor while preserving quasi-dynamic variable metadata on older releases. Added lowering for JAX 0.11.2’s lax.one_minus_square as (1 + x) * (1 - x), retaining precision near |x| = 1 and fixing the generated CI cases for tanh, acos, asin, and atanh gradients. The scalar constant follows the input dtype, including in float32 lax.scan bodies when double-precision export is enabled.

• Improved cosine accuracy across dtypes: Switched jnp.cos below float64 to native ONNX Cos, avoiding phase-shift error in float32 rotary embeddings. Float64 jnp.cos and lax.cos now use 1 - 2 * Sin(x / 2)², avoiding the large-angle error of Sin(x + π/2) while retaining compatibility with ONNX Runtime 1.24.1, whose CPU kernels lack double-precision Cos. Regression tests cover large arguments, cosine zeros, and nonfinite values with runtime optimizations enabled and disabled, plus scalar and symbolic shapes and float32 rotary-embedding angles.

• Preserved FP16 matrix-product types: Fixed ONNX output type annotations for jnp.matmul and jnp.dot with FP16 operands, and added operand casts when an explicit preferred_element_type requests wider accumulation. Added CPU runtime checks for matrix products, NNX Linear, and bias-plus-GELU with static and dynamic batches and runtime optimizations enabled and disabled. These checks validate FP16 outputs while allowing ONNX Runtime to use internal FP32 fallback for FP16 operations.

• Made decoder normalization tradeoffs reproducible: Added a CPU benchmark comparing auto and prefer_native on identical tiny-decoder weights and inputs. The benchmark covers static and symbolic exports, JAX parity, optimized operator counts, and runtime and hardware metadata. Updated the native-kernel description for ONNX Runtime 1.30’s AVX2 implementation while retaining the fixed accuracy gates.

• Strengthened numerical validation and CI reproducibility: Added CPU accuracy checks on Python 3.12/JAX 0.10.2 and Python 3.13/JAX 0.11.2, with fixed bounds for selected LayerNorm and cosine cases. Expanded BF16 export capability coverage for explicit Linen LayerNorm and added static/dynamic NNX decoder examples using explicit auto graphs. Enabled CI for documentation-only changes, version-tag pushes, and manual runs, and recorded the tested commit, dependency versions, runner image, and CPU hardware in each accuracy job’s summary.

Commits

• d0534a1 supported components update
• f7766a1 fix: last version 0.16.1
• 38f4aaa release preparation
• 72e3d2c fix: improve cosine and FP16 exports for ONNX Runtime 1.30 (#264)
• 6df852a docs(ci): clarify normalization accuracy scope and enable release checks
• 1cc743f fix: lower JAX one_minus_square to restore gradient CI tests
• b595d3a dependencies update
• 9a2f6ff Fix transposed-conv stride loss and opset-invalid Gelu exports (#261)
• eb4e250 Repair torch 2.13.0 lock update (#259)
• c19e77e Repair mkdocstrings 1.0.6 lock update (#258)
• Additional commits viewable in compare view

[coderabbitai]

Important

Review skipped

Bot user detected.

To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

• Configuration used: defaults
• Review profile: CHILL
• Plan: Advanced
• Run ID: 64d799c8-c799-4ded-886c-3bef09c6574e

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

• Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[AlexanderFengler]

@dependabot rebase