Migrate to new translation function (#18)

* feat: optimisiing comp graph

* feat: graph tranlsation

* Feat/pendulum rollout example (#9)

* added pendulum rollout as example

* fixed README

* fix: rebase

* Update README.md

Added hyperlink for landing

* fix: update branch for colab benchmark

* added demo notebook (#13)

* added demo notebook

* minor readme change

* fix: rebase

* fix: rebase and fix tests

* fix: pre-commit

* feat: graph expansion

* fix: benchmarks for more powers

* fix: disable graph compression

* fix: densify structural zeros

* fix: adaprive dimensionality #16

* fix: translate as graph_translate

* feat: test expand, docs

* feat: test examples

* fix: exclude running examples from workflow

* fix: remove ast debug files

* fix: organize imports

* fix: ops fix

* fix: update plots for graph translation

* fix: add plots to website

---------

Co-authored-by: Simeon Nedelchev <[email protected]>
Co-authored-by: Lev Kozlov <[email protected]>

Author: 3 people

.github/workflows/build.yaml

@@ -38,7 +38,7 @@ jobs:
           python -m pip install --upgrade pip
           pip install -e .[dev]
       - name: Run tests
-        run: python -m pytest
+        run: python -m pytest --ignore=tests/test_examples.py
 
   deploy:
     runs-on: ubuntu-latest

benchmarks/cuda_benchmark_results.npz

@@ -17,7 +17,7 @@
 class Paths:
     FUNCTIONS_DIR = os.path.join(cur_dir, "data")
     OUTPUT_DIR = cur_dir
-    RUN_CUSADI = False
+    RUN_CUSADI = True
 
 
 class ColabPaths:
@@ -43,11 +43,11 @@ class ColabPaths:
 
     torch.manual_seed(0)
 
-N_ENVS_SWEEP = [2**i for i in range(21)]
-N_EVALS = 20
+N_ENVS_SWEEP = [2**i for i in range(15)]
+N_EVALS = 10
 
 # Load functions for CUDA benchmarking
-fn_files = ["fn_1e1.casadi", "fn_1e2.casadi"]
+fn_files = ["fn_1e1.casadi", "fn_1e2.casadi", "fn_1e3.casadi", "fn_1e4.casadi"]
 benchmark_fns = [Function.load(os.path.join(PathsProvider.FUNCTIONS_DIR, fn)) for fn in fn_files]
 
 

benchmarks/jaxadi_benchmark_results.npz

@@ -166,6 +166,12 @@ <h2 class="title is-3">Benchmarks</h2>
                     <div class="item item-chair-tp">
                         <img src="./static/images/compare_1e2_bar.png"/>
                     </div>
+                    <div class="item item-chair-tp">
+                        <img src="./static/images/compare_1e3_bar.png"/>
+                    </div>
+                    <div class="item item-chair-tp">
+                        <img src="./static/images/compare_1e4_bar.png"/>
+                    </div>
                     <div class="item item-steve">
                         <img src="./static/images/speedup_ratio.png"/>
                     </div>