Keras 3 determinism behavior across JAX and PyTorch backends

Determinism in Keras 3 Multi-Backend Environments

No, keras.config.enable_determinism() does not automatically synchronize or trigger the low-level random number generator (RNG) settings for JAX and PyTorch. While this global flag manages high-level Keras configurations and Python/NumPy seeds, it does not replace the need for backend-specific seeding mechanisms.

Backend-Specific Requirements

To achieve full reproducibility, you must manually integrate the following settings alongside the Keras determinism flag:

• PyTorch: You must call torch.manual_seed(seed). If using NVIDIA GPUs, you must also set torch.backends.cudnn.deterministic = True and torch.backends.cudnn.benchmark = False to disable non-deterministic CUDA kernels.
• JAX: JAX uses a stateless PRNG system. You must explicitly create a PRNG key using jax.random.PRNGKey(seed) and ensure that the same key sequence is passed to any random operations. Keras does not internally manage the JAX PRNG key state.

Recommended Verification Workflow

To verify that weight updates are identical across backends, follow this scoped test procedure:

1. Initialize: Set a fixed seed (e.g., 42) and apply keras.config.enable_determinism(True). Apply the backend-specific seeds mentioned above.
2. Standardize: Build the same model architecture and compile it with the same optimizer and loss function across both backends.
3. Execute: Train on a small, fixed dataset for a limited number of steps (e.g., 2 epochs).
4. Compare: Extract the model weights and compare them using numpy.allclose with a floating-point tolerance (e.g., atol=1e-6).

import numpy as np
# Example verification check
np.allclose(weights_jax, weights_pytorch, atol=1e-6)

Assumptions and Limitations

This guidance assumes the use of Keras 3.x. Be aware that enabling deterministic operations often disables optimized parallel algorithms, which may significantly increase training time on GPUs. Additionally, slight numerical divergences may still occur due to differences in how different backends handle floating-point reduction orders.

Diagnostic Detail Needed: Are you utilizing any custom layers with stochastic behavior (e.g., custom Dropout implementations) that bypass standard Keras ops?