Dynamism in StableHLO

The current state of dynamism is more formally spelled out in the Dynamism RFC, this page will provide a high level overview of the RFC and discuss important APIs and tooling for interacting with dynamic programs.

Dynamism Terminology & Support Overview

First, to cover a few terms that will appear in this doc, as well as a brief intro to their support in StableHLO:

Dynamic dimensions

Dynamic dimensions refers to any dimension whose dimension size is unknown. In StableHLO we represent dynamic dimensions using ?, i.e. tensor<16x?xf32>.

Bounded dynamism

Bounded dynamism refers to a dynamic dimension whose value has a known upper bound. Generally this is useful for padding the tensor during execution. In StableHLO we represent bounded dynamism using #stablehlo.bounds as a tensor encoding, i.e. a rank-2 tensor with one dynamic dimension bounded at 16 and the other without a bound can be represented as tensor<?x?xf32, #stablehlo.bounds<16, ?>>.

StableHLO is able to represent bounded dynamism, but there is limited framework support, originating in TensorFlow, and with some support in PyTorch/XLA.

Unbounded dynamism

Unbounded dynamism as the name implies refers to a dynamic dimension with no known bound on the size. This type of dynamism is very common in StableHLO, with JAX, PyTorch/XLA, and TF support, often used for exporting models with dynamic batch size or sequence length.

In StableHLO we simply elide the bounds encoding for this form of dynamism, i.e. tensor<?x?xf32>.

Shape polymorphism

Shape polymorphism is a term we've inherited from JAX.

There are two key implications to shape polymorphism:

1. All dynamism in the program traces back to its input arguments.
2. All dynamism pertains to tensor shapes only, i.e. not data-dependent.

With these two rules, once the static shapes of a program are known, we are able to take a dynamic program and fully refine it into a static program for compilation (see "Compiler passes for refining dynamic programs").

Generally shape polymorphism uses unbounded dynamism, if known argument shapes can lead to a fully static program, there isn't a need to guess on how to bound the values.

Data-dependent dynamism

Data-dependent dynamism refers to dynamic dimensions sizes that pertain to the data inside a tensor. The canonical example is a nonzeros function which returns the indices of all elements that are 0 in a tensor value. The shape cannot be known without evaluating the data, but it can often be compiled using bounded dynamism, spending extra memory on the potential output tensor size.

Many data-dependent dynamic ops can be modeled using bounded dynamism, where an upper bound on a tensor size is specified, and hardware generally will implement this via tensor padding. Today there is some support for data-dependent dynamism in PyTorch/XLA and TensorFlow, but JAX does not currently trace operations which lead to data dependent dynamism.

Exporting programs with dynamic dimensions

See our StableHLO tutorials for information on how to export programs with dynamic batch sizes or sequence lengths:

• JAX Tutorial > Export with Dynamic Batch Size
• PyTorch/XLA Tutorial > Export with Dynamic Batch Size

Compiler passes for refining dynamic programs

Remove dynamism pass pipeline

There are a few useful passes for refining shapes, conveniently they are all bundled in a pass pipeline createStablehloRemoveDynamismPipeline:

void createStablehloRemoveDynamismPipeline(OpPassManager &pm,
                                           TypeRange refinedTypes);

Individual passes for refining dynamism

Individually, the passes that tend to be useful for shape refinement are:

• stablehlo-refine-arguments to replace input arguments with concrete tensor types.
• stablehlo-refine-shapes to propagate the new input argument shape information throughout the entire program.
• stablehlo-canonicalize-dynamism to replace dynamic ops with their static variants.
• stablehlo-check-shape-assertions to check and remove shape assertions custom calls.

See linked documentation for up-to-date information and examples.

Example: How is dynamism useful, and how can I use it?

Dynamism has lots of uses, here we'll mainly focus on the common use case for Shape Polymorphism - creating a flexible exported model representation, generally used to represent dynamic batch size or sequence length.

Static add_one model

We'll use the following simple add_one model to demonstrate this: