ADTests

This repository is used to tabulate the current performance of various automatic differentiation (AD) backends with Turing.jl models.

Let me see the results!

https://turinglang.org/ADTests

I want to add more AD backends!

You can modify the list of AD types in main.jl.

I want to add more models!

You can modify the list of models by adding a new file inside the models directory. This file should contain the model definition, and call the @register macro to register the model with the ADTests package. See the existing files in that directory for examples.

To make sure that the definition is included in the final website, you will have to make sure that the filename is the same as the model name (i.e. a model @model function f() is in models/f.jl).

I want to edit the website!

The website is a small Svelte app in the web directory.

To build this website locally, you will need to:

1. Install pnpm if you don't already have it. (npm is fine too, just replace pnpm with npm run in the commands below)

2. Download the JSON files from the gh-pages branch, and place them in the web/src/data directory. Currently, there are three JSON files: adtests.json, manifest.json, and model_definitions.json. These represent the latest results from running the AD tests on CI.

3. cd web

4. pnpm install

5. pnpm dev

6. Open http://localhost:5173 in your browser.

I want to see the HTML generated by a PR!

The latest workflow run across all PRs will be published to https://turinglang.org/ADTests/pr.

This is a bit messy, but works for now on the assumption that there aren't many PRs being worked on simultaneously.

What's going on?

The workflow is the most complicated part of this repository. This section attempts to explain it from the 'bottom up'; if you prefer a 'top down' approach start by looking at the GitHub Action workflow, .github/workflows/test.yml.

Under the hood, the main thing that actually runs the AD tests / benchmarks is main.jl. You can run julia --project=. main.jl and it will print some usage information. However, it is the Python script ad.py that controls how this Julia script is called.

Fundamentally, the idea is that we want to loop over every combination of model and adtype and test it. However, because GitHub Actions limits jobs to 6 hours, it is impractical to run every combination in the same job. What we do is to run one job per model. This is accomplished by first storing the names of the models and adtypes in the $GITHUB_OUTPUT variable using python ad.py setup, which can then be read by the next job.

The next job is a CI matrix split by model name; each of the sub-jobs invokes python ad.py run --model {model_name}, which loops over each adtype and calls the Julia script with the model name and adtype as arguments. The purpose of having this Python -> Julia setup (as opposed to doing the looping inside Julia itself) is to guard against the Julia process crashing, which can happen sporadically with Enzyme. If the Julia process successfully finishes, it will print the result which is picked up by the Python script; if it crashes, we just record the result as 'error'.

Finally, the results are collated and sent to the final job in the workflow, which is python ad.py html. This bit of the Python script is responsible for generating the three JSON files which the web app uses. (Fun fact: collating these results is also somewhat involved because we can't just write to $GITHUB_OUTPUT; it turns out that output from different jobs in a matrix will override each other, so the output can't share the same key, and there's no way to dynamically specify the output key. Thankfully, there is an existing action which is designed to get around this problem by uploading artefacts instead of using $GITHUB_OUTPUT.)

Overall, what this means is that the entire table can be generated in around 10 minutes (longer if you need to install + precompile dependencies, but on GitHub Actions dependencies will for the most part have been cached).

Can I run this locally?

If you just want to run one specific combination of model and adtype, you can run julia --project=. main.jl --run <model> <adtype>.

If you want to run one model with all adtypes, you can run uv run ad.py run --model <model_name>.

You probably don't want to run all models with all adtypes, as that takes a really long time.

If you just want to build the website using results that were obtained from CI, that's described in the section above.