Overview

TensorFlow 2.16, released on April 11, 2024, uses Keras 3 by default and supports JAX as a backend.

Main Features

Keras 3 by default

Keras 3 is the default backend, providing a multi-framework API compatible with TensorFlow, JAX, and PyTorch.

python 
import keras

model = keras.Sequential([
    keras.layers.Dense(128, activation='relu'),
    keras.layers.Dense(10, activation='softmax'),
])
model.compile(optimizer='adam', loss='sparse_categorical_crossentropy')
print(keras.backend.backend())  # tensorflow

JAX support

Keras 3 allows using JAX as a backend to benefit from XLA compilation.

python 
# KERAS_BACKEND=jax python script.py
import os
os.environ['KERAS_BACKEND'] = 'jax'

import keras
print(keras.backend.backend())  # jax

Sources