Cross-framework model conversion, especially from TensorFlow to JAX, is challenging due to behavioral differences in layers, defaults, and state handling. For researchers building translation tools or evaluating framework fidelity, there was no available dataset with:

• Equivalent TF and JAX model implementations
• Deep layer-level mappings capturing configuration, structure, and edge cases
• Consistent, traceable annotations to enable numerical and functional comparison

The client needed a benchmark-grade dataset to serve as a ground truth for both manual and machine-assisted translation systems.

Turing built a benchmark dataset to evaluate framework interoperability, designed for traceability, semantic alignment, and expert verification.

Model sourcing

• Sourced more than 30 public TensorFlow models with permissive licenses and architectural diversity
• Included generative models, transformers, reinforcement learning models, graph neural networks, computer vision models, and recommendation systems

Manual translation

• Rewrote each model in JAX/Flax, ensuring semantic alignment with the original version
• Verified that each JAX model ran without trivial errors and matched expected layer behavior

Layer-level annotation

• Mapped each TF layer to its JAX equivalent (tf.keras.layers.Dense to flax.linen.Dense)
• Documented input and output shapes, parameter configurations, and initializer differences
• Captured behavior notes on state handling such as BatchNorm train vs. eval modes, numerical sensitivities, and conversion challenges
• Logged precise mappings of weight and bias names, shapes, and roles

Unit testing

• Wrote unit tests for every model pair to validate consistency in layer counts and total parameter counts between TF and JAX implementations
• Enabled rapid validation of translation fidelity at the architectural level

Numerical equivalence testing

• Conducted forward-pass equivalence tests between each original TensorFlow model and its JAX translation
• Verified that output arrays from both models matched within a strict tolerance threshold (≤ 0.01%)
• Ensured that each translation preserved the original model’s numerical behavior to support benchmark validity

Data packaging

• Delivered translated models, source code, and detailed annotations in a structured format
• Aligned all assets to the client’s interoperability benchmarking infrastructure

• Delivered over 500 deep layer annotations across more than 30 model pairs
• Ensured ≤ 0.01% numerical equivalence tolerance for every translated model
• Enabled traceable numerical fidelity comparisons between TF and JAX implementations
• Structured all data for direct use in benchmarking framework equivalence
• Identified edge-case behaviors and inconsistencies in stateful and compound layers

This dataset serves as a trusted baseline for evaluating framework interoperability at a layer level. With it, the client can:

• Benchmark genAI translation tools on real TF-to-JAX fidelity
• Identify framework-specific behavior drift in complex models
• Extend toward automated verification of layer behavior across deep learning libraries
• Train or validate future translation models using task-aligned, expert-labeled samples