StableHLO Model Exports

Exports of popular models in StableHLO.

Exports with large literals (weights) elided are available in the exports/ folder.

Currently this includes:

• GDM AlphaFold
• GDM Searchless Chess
• FlaxResNet50
• hf_BERT
• More to come

I recommend a fresh venv for exporting, this has been tested only using py3.11 but py3.10 is likely to work as well.

python3 -m venv venv
source venv/bin/activate
(cd models/ && ./setup.sh)

To export all models use the all flag, this is the default so it can be omit.

export PYTHONPATH="$(pwd)/models:$(pwd)/models/gdm_searchless_chess"
python models/export.py --models=all

A list of models can be specified using the models flag if needed:

python models/export.py --models=searchless_chess_9m

All exports can be viewed in the exports/ folder, the .mlir files have large constants elided for readability. All .mlir.bc files have the constants embedded.

Note: Some of the files with many large constants do not have .mlir.bc files since they don't fit in git well. For these files you will need to run the export scripts, or open a ticket with the request end we can figure out where to host these.

To view the .mlir.bc file, there are two options:

If you have access to stablehlo-opt and stablehlo-translate from building the StableHLO repository:

stablehlo-translate --deserialize alphafold.mlir.bc

or

stablehlo-opt --pass-pipeline='builtin.module(stablehlo-deserialize)' alphafold.mlir.bc

On Linux we have ready-to-use python bindings for StableHLO published nightly:

pip install stablehlo -f https://github.com/openxla/stablehlo/releases/expanded_assets/dev-wheels

Once installed the following script can be used to deserialize and print the IR:

from mlir.dialects import stablehlo
from mlir.ir import Context, Location

filepath="/path/to/stablehlo-exports/exports/alphafold.mlir.bc"
with open(filepath, 'rb') as file:
  bytecode = file.read()

with Context() as ctx:
  module = stablehlo.deserialize_portable_artifact(ctx, bytecode)

print(module)