[ICML'25] Price of Freedom: Exploring Tradeoffs between Expressivity and Computational Efficiency in Equivariant Tensor Products

The official implementation of the Price of Freedom: Exploring Tradeoffs between Expressivity and Computational Efficiency in Equivariant Tensor Products, published at ICML 2025.

Installation

We recommend using uv:

uv sync
source .venv/bin/activate
uv pip install -e .

Alternatively, you can install the dependencies manually:

python -m venv .venv
source .venv/bin/activate
uv pip install -r requirements.in

Tetris Experiments

bash shell/run_tetris_experiments.sh

3BPA experiments

Checkout experiments/Gaunt-Tensor-Product-S2Grid. Refer to the original code for reproducing the experiments

Plotting

experiments/Gaunt-Tensor-Product-S2Grid/analyze_results.ipynb

Benchmarking experiments

CPU

bash shell/run_cpu_timing.sh

GPU

bash shell/run_gpu_timing.sh

Nsight Compute installation needed. Make sure the GPU counters are enabled.

bash shell/run_nsight_profiling.sh

Plotting

notebooks/benchmarking_viz