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>

benchmarks/run_benchmark.py

@@ -8,7 +8,7 @@
 
 import casadi as cs
 
-from jaxadi import translate
+from jaxadi import graph_translate as translate
 
 # define input variables for the function
 x = cs.SX.sym("x", 3)

benchmarks/visualize.ipynb

@@ -1,17 +1,20 @@
 import casadi as cs
 
-from jaxadi import lower, translate, declare
+from jaxadi import lower, declare
+from jaxadi import graph_translate as translate
 
 # define input variables for the function
-x = cs.SX.sym("x", 10, 10)
-y = cs.SX.sym("y", 10, 10)
-casadi_fn = cs.Function("myfunc", [x, y], [x @ y])
+x = cs.SX.sym("x", 2, 1)
+y = cs.SX.sym("y", 2, 1)
+casadi_fn = cs.Function("myfunc", [x, y], [x.T @ y])
 
 print("Signature of the CasADi function:")
 print(casadi_fn)
 
 # define jax function from casadi one
 jax_fn = declare(translate(casadi_fn))
 
+print(translate(casadi_fn))
+
 print("Lowered JAX function:")
 print(lower(jax_fn, casadi_fn).as_text())

docs/index.html

@@ -16,7 +16,8 @@
 import pinocchio as pin
 import pinocchio.casadi as cpin
 from robot_descriptions.panda_description import URDF_PATH
-from jaxadi import convert, translate
+from jaxadi import convert
+from jaxadi import graph_translate as translate
 
 # Load the Panda robot model
 model = pin.buildModelFromUrdf(URDF_PATH)

docs/static/images/compare_1e1_bar.png

@@ -1,4 +1,5 @@
 from ._compile import lower
 from ._convert import convert
-from ._translate import translate
 from ._declare import declare
+from ._expand import translate as expand_translate
+from ._graph import translate as graph_translate

docs/static/images/compare_1e2_bar.png

@@ -1,13 +1,15 @@
-from casadi import Function
-from typing import Any
 from collections.abc import Callable
+from typing import Any
+
+from casadi import Function
 
-from ._declare import declare
-from ._translate import translate
 from ._compile import compile as compile_fn
+from ._declare import declare
+from ._graph import translate as graph_translate
+from ._preprocess import densify
 
 
-def convert(casadi_fn: Function, compile=False, num_threads=1) -> Callable[..., Any]:
+def convert(casadi_fn: Function, translate=None, compile=False) -> Callable[..., Any]:
     """
     Convert given casadi function into python
     callable based on JAX backend, optionally
@@ -17,7 +19,12 @@ def convert(casadi_fn: Function, compile=False, num_threads=1) -> Callable[...,
     :param compile (bool): Whether to AOT compile function
     :return (Callable[..., Any]): Resulting python function
     """
-    jax_str = translate(casadi_fn, num_threads=num_threads)
+    if translate is None:
+        translate = graph_translate
+
+    casadi_fn = densify(casadi_fn)
+
+    jax_str = translate(casadi_fn)
     jax_fn = declare(jax_str)
 
     if compile: