Overview

TensorFlow 2.15, released on January 18, 2024, prepares the migration to Keras 3 and improves multi-backend compatibility.

Main Features

Keras 3 preparation

TensorFlow 2.15 ships Keras 2 by default but lets you test Keras 3 which supports JAX and PyTorch as alternative backends.

python 
import os
os.environ['TF_USE_LEGACY_KERAS'] = '0'  # Keras 3

import keras
model = keras.Sequential([
    keras.layers.Dense(64, activation='relu'),
    keras.layers.Dense(10, activation='softmax'),
])
model.compile(optimizer='adam', loss='categorical_crossentropy')
print(f'Backend: {keras.backend.backend()}')

XLA improvements

The XLA compiler improves JIT compilation performance and reduces memory footprint of compiled models.

python 
import tensorflow as tf

@tf.function(jit_compile=True)
def train_step(x, y):
    with tf.GradientTape() as tape:
        pred = model(x, training=True)
        loss = loss_fn(pred, y)
    grads = tape.gradient(loss, model.trainable_variables)
    optimizer.apply_gradients(zip(grads, model.trainable_variables))
    return loss

Sources