Overview

TensorFlow 2.8, released on February 4, 2022, introduces async dispatch and the DTensor API for distributed computing.

Main Features

Async dispatch

Async dispatch improves performance by sending operations to the GPU without waiting for completion, enabling better compute/transfer overlap.

python 
import tensorflow as tf

# Async dispatch is enabled by default
model = tf.keras.Sequential([
    tf.keras.layers.Dense(128, activation='relu'),
    tf.keras.layers.Dense(10, activation='softmax'),
])
model.compile(optimizer='adam', loss='sparse_categorical_crossentropy')

# Training benefits from async dispatch
# model.fit(x_train, y_train, epochs=5)

DTensor

The DTensor API distributes tensors and computations across multiple devices (GPU/TPU) transparently, simplifying data and model parallelism.

python 
import tensorflow as tf
from tensorflow.experimental import dtensor

# Create a compute mesh
mesh = dtensor.create_mesh([("batch", 2)], devices=['CPU:0', 'CPU:1'])
layout = dtensor.Layout(['batch', dtensor.UNSHARDED], mesh)

# Distributed tensor on the mesh
tensor = dtensor.call_with_layout(
    tf.ones, layout, shape=(4, 3)
)
print(tensor.shape)  # (4, 3)

Sources