Keras 3 backend transitions and gradient computation parity

Gradient Parity Across Backends

Keras 3 does not guarantee identical gradient values across different backends. While keras.ops provides a unified API for high-level portability, the underlying automatic differentiation engines (JAX, PyTorch, and TensorFlow) utilize different kernels, floating-point accumulation strategies, and XLA compilation paths. Consequently, you should expect numerical divergence rather than bit-wise parity.

Factors Contributing to Divergence

• Numerical Precision: Differences in how each backend handles floating-point rounding and epsilon values in operations like LayerNormalization or Adam optimization.
• Autodiff Implementations: Each backend implements the chain rule and gradient accumulation differently at the primitive level.
• Hardware Kernels: Backend-specific optimizations for GPU/TPU may result in slight variations in the order of operations, affecting final gradient sums.

Handling Optimizer State for Continuity

Keras 3 does not provide an automated, seamless "hot-swap" for optimizer states between backends. Because optimizer states (such as momentum buffers in SGD or first/second moments in Adam) are stored as backend-specific tensors, they must be manually transitioned to ensure training continuity.

Steps for State Migration

1. Extract Weights and State: Convert the current weights and optimizer state tensors into a neutral format, such as NumPy arrays.
2. Switch Backend: Change the KERAS_BACKEND environment variable to the target framework.
3. Reconstruct Model: Re-instantiate the model and optimizer in the new backend.
4. Inject State: Load the NumPy arrays back into the new backend's tensor format and assign them to the model weights and optimizer variables.

Verification of Parity

To quantify the divergence when switching backends, you can perform a single-step gradient check using the following approach:

import numpy as np
import keras

# Assuming model and input are defined
with keras.GradientTape() as tape:
    loss = model(input_data)

grads = tape.gradient(loss, model.trainable_variables)
# Convert to numpy and compare across backends
# np.allclose(grads_jax, grads_pytorch, atol=1e-5)

Missing Diagnostic Detail: Are you utilizing mixed-precision training (e.g., float16 or bfloat16)? The choice of precision significantly impacts the magnitude of divergence between backends.