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:

1. Adding a new call to @include_model {category_heading} {model_name} in main.jl. Both category_heading and model_name should be strings.

   category_heading is used to determine which table the model appears under on the website. This should be self-explanatory if you look at the current website.

2. Adding a new file, models/{model_name}.jl.

   The basic structure of this file should look like this, where model_name is replaced accordingly:

   #=
   (1) You can add any explanatory comments here if necessary
   =#
   
   # (2) Imports if necessary
   using MyOtherPackage: some_function
   
   # (3) Define data if necessary
   data = ...
   
   # (4) Define the model
   @model function model_name(data, ...)
       # Define your model here
       ...
   end
   
   # (5) Instantiate the model
   model = model_name(data, ...)

   (1) Description

   Ideally, model_name would be self-explanatory, i.e. it would serve to illustrate exactly one feature and the name would indicate this. However, if necessary, you can add further explanatory comments inside the model definition file.

   (2) Dependencies

   Inside this file, you do not need to call using Turing or any of the AD backends. (This also means you do not need to import anything that Turing re-exports, such as distributions.)

   However, you will have to make sure to import any other packages that your model requires. (If this package is not already present in the project environment, you will also have to add it to Project.toml.)

   (3) Data definition

   Each file in models/ is evaluated within its own module, so you can declare data variables, etc. without worrying about name clashes.

   (4) Model definition

   Models can be defined as usual with @model function model_name(...).

   (5) Model instantiation

   The last line in the file should be the creation of the Turing model object using model = model_name(...). (It is mandatory for the model object to be called model.)

Important

Note that for CI to run properly, model_name must be consistent between the following:

• The name of the model itself i.e. @model function model_name(...)
• The filename i.e. models/model_name.jl
• The name of the model in main.jl i.e. @include_model "Category Heading" "model_name"

(This setup does admittedly feel a bit fragile. Unfortunately I could not find a simpler way to get all the components (Julia, Python, web app) to work together in an automated fashion. Hopefully it is a small price to pay for the ability to just add a new model and have it be automatically included on the website.)

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.