TF and JAX — Google's Two Frameworks

Why Google maintains two frameworks (and what it means for you)

Google maintains two major ML frameworks: TensorFlow (2015) and JAX (2018). Understanding the difference is essential context for anyone working in the TF ecosystem in 2026.

Dimension	TensorFlow	JAX
Paradigm	OO, stateful	Functional, pure functions
Model state	Object attributes	Explicit pytrees
Key ops	tf.function, GradientTape	jit, grad, vmap, pmap
Random numbers	Global RNG	Explicit PRNG key splitting
Debug experience	Good (eager default)	Hard under jit
TPU performance	Good (via XLA)	Best-in-class
Production deployment	Excellent (Serving, TFLite, TF.js)	Hard (needs jax2tf bridge)

Google's frontier LLMs (Gemini, PaLM 2, Gemma) are trained in JAX. TensorFlow remains dominant for production serving and mobile deployment. The interoperability bridge is jax2tf — convert JAX functions to TF for TF Serving / TFLite deployment.

The bigger picture: TF, JAX, and PyTorch are converging on StableHLO as a common intermediate representation for XLA hardware. Long term, the framework you write in matters less than the deployment target — Keras 3 is the most visible part of this convergence.

Code

jax2tf — JAX 함수를 TF SavedModel로·python

from jax.experimental import jax2tf
import jax.numpy as jnp
import tensorflow as tf

def jax_model(x):
    return jnp.sin(jnp.cos(x))

# Convert JAX function to TF (for TF Serving / TFLite)
tf_model = jax2tf.convert(jax_model)
result = tf_model(tf.constant([1.0, 2.0]))

# Export as TF SavedModel from JAX training
tf_func = tf.function(tf_model, autograph=False)
module   = tf.Module()
module.f = tf_func
tf.saved_model.save(module, "/tmp/jax_as_tf_saved_model")

TF and JAX — Google's Two Frameworks

Why Google maintains two frameworks (and what it means for you)

Code

Progress

Comments 0