Skip to content

Commit 537aad6

Browse files
mattephilvjonok
andauthored
Benchmarks + landing page (#10)
* fix: solve problem with different shapes in input * feat: start landing * fix: favicon * feat: add benchmark util func * fix: try cuda first * fix: merge new translate * fix: outputs * fix: try with memory cleaning * fix: change order of jit vmap * fix: increase batch * fix: prints * fix: no tqdm * feat: increase the batch * fix: reduce * fix: no tqdm * feat: add visualize demo * feat: add notebook from colab * fix: plots * feat: add patches and update readme * fix: benchmarks on landing * fix: landing * fix: landing * fix: expand benchmarking readme * fix: bump version --------- Co-authored-by: Lev Kozlov <[email protected]>
1 parent 56ee252 commit 537aad6

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

41 files changed

+4257
-100
lines changed

CONTRIBUTING.md

Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,56 @@
1+
# Contributing to JaxADi
2+
3+
We're thrilled that you're interested in contributing to JaxADi! This document outlines the process for contributing to this project.
4+
5+
## How to Contribute
6+
7+
There are many ways to contribute to JaxADi:
8+
9+
1. Reporting bugs
10+
2. Suggesting enhancements
11+
3. Writing documentation
12+
4. Submitting code changes
13+
14+
### Reporting Bugs
15+
16+
1. Check the [issue tracker](https://github.com/based-robotics/jaxadi/issues) to see if the bug has already been reported.
17+
2. If not, create a new issue. Provide a clear title and description, as much relevant information as possible, and a code sample or executable test case demonstrating the bug.
18+
19+
### Suggesting Enhancements
20+
21+
1. Check the [issue tracker](https://github.com/based-robotics/jaxadi/issues) to see if the enhancement has already been suggested.
22+
2. If not, create a new issue. Clearly describe the enhancement, why it would be useful, and any potential drawbacks.
23+
24+
### Writing Documentation
25+
26+
Good documentation is crucial. If you notice any part of our documentation that could be improved or expanded, please let us know or submit a pull request with your suggested changes.
27+
28+
### Submitting Code Changes
29+
30+
1. Fork the repository.
31+
2. Create a new branch for your changes.
32+
3. Make your changes in your branch.
33+
4. Add or update tests as necessary.
34+
5. Ensure the test suite passes.
35+
6. Update the documentation as needed.
36+
7. Push your branch and submit a pull request.
37+
38+
## Pull Request Process
39+
40+
1. Ensure your code follows the project's style guidelines.
41+
2. Update the README.md or relevant documentation with details of changes, if applicable.
42+
3. Add tests for your changes and ensure all tests pass.
43+
4. Your pull request will be reviewed by the maintainers. They may suggest changes or improvements.
44+
5. Once approved, your pull request will be merged.
45+
46+
<!-- TODO:
47+
48+
## Style Guidelines
49+
50+
We use Ruff to enforce our code style and linting rules....
51+
```bash
52+
pip3 install pre-commit
53+
```
54+
```bash
55+
pre-commit install
56+
``` -->

README.md

Lines changed: 36 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -1,32 +1,33 @@
11
<!-- # JaxADi -->
2-
<!-- TODO: ADD PATCHES -->
2+
3+
[![CI](https://img.shields.io/github/actions/workflow/status/based-robotics/jaxadi/build.yaml?branch=master)](https://github.com/based-robotics/jaxadi/actions)
4+
[![PyPI version](https://img.shields.io/pypi/v/jaxadi?color=blue)](https://pypi.org/project/jaxadi/)
5+
[![PyPI downloads](https://img.shields.io/pypi/dm/jaxadi?color=blue)](https://pypistats.org/packages/jaxadi)
36

47
<p align="center">
58
<!-- Placeholder for a cool logo -->
69
<img src="https://github.com/based-robotics/jaxadi/blob/master/_assets/_logo.png?raw=true" alt="JAXADI Logo" width="400"/>
710
</p>
811

9-
1012
**JaxADi** is a powerful Python library designed to bridge the gap between `casadi.Function` and JAX-compatible functions. By leveraging the strengths of both CasADi and JAX, JAXADI opens up exciting opportunities for building highly efficient, batchable code that can be executed seamlessly across CPUs, GPUs, and TPUs.
1113

1214
JAXADI can be particularly useful in scenarios involving:
1315

1416
- Robotics simulations
1517
- Optimal control problems
1618
- Machine learning models with complex dynamics
17-
- Large-scale numerical optimizations
18-
1919

2020
## Installation
2121

2222
You can install JAXADI using pip:
23+
2324
<!-- Change once it will be realeased -->
2425

2526
```bash
2627
pip install jaxadi
2728
```
2829

29-
For a complete environment setup, we recommend using Conda/Mamba:
30+
For a complete environment setup for examples, we recommend using Conda/Mamba:
3031

3132
```bash
3233
mamba env create -f environment.yml
@@ -63,32 +64,53 @@ output = jax_fn(input_x, input_y)
6364

6465
```
6566

67+
<strong>Note:</strong> For now translation does not support functions with very
68+
large number of operations, due to the translation implementation. Secret component of
69+
translation is work-tree expansion, which might lead to large overhead in number of
70+
symbols. We are working on finding the compromise in both speed and extensive
71+
functions support.
72+
6673
## Examples
6774

6875
JAXADI comes with several examples to help you get started:
6976

70-
1. [Basic Translation](examples/00_translate.py): Learn how to translate CasADi functions to JAX.
77+
1. [Basic Translation](examples/00_translate.py): Learn how to translate CasADi functions to JAX.
7178

72-
2. [Lowering Operations](examples/01_lower.py): Understand the lowering process in JaxADi.
79+
2. [Lowering Operations](examples/01_lower.py): Understand the lowering process in JaxADi.
7380

74-
3. [Function Conversion](examples/02_convert.py): See how to fully convert CasADi functions to JAX.
81+
3. [Function Conversion](examples/02_convert.py): See how to fully convert CasADi functions to JAX.
7582

76-
4. [Pendulum Rollout](examples/03_pendulum_rollout.py): Batched rollout of the nonlinear passive nonlinear pendulum
83+
4. [Pendulum Rollout](examples/03_pendulum_rollout.py): Batched rollout of the nonlinear passive nonlinear pendulum
7784

78-
5. [Pinocchio Integration](examples/04_pinocchio.py): Explore how to convert Pinocchio-based CasADi functions to JAX.
85+
5. [Pinocchio Integration](examples/04_pinocchio.py): Explore how to convert Pinocchio-based CasADi functions to JAX.
7986

8087
6. [MJX Comparison](examples/05_mjx.py): Compare the transformed Pinnocchio forward kinematics with one provided by Mujoco MJX
8188

8289
> **Note**: To run the Pinocchio and MJX examples, ensure you have them properly installed in your environment.
8390
91+
## Performance Benchmarks
8492

85-
<!-- ## Performance Benchmarks
93+
![speedup](https://github.com/based-robotics/jaxadi/blob/master/docs/static/images/speedup_ratio.png?raw=true)
8694

87-
(Consider adding a section about performance comparisons between CasADi and JAXADI-translated functions) -->
95+
The process of benchmarking and evaluating the performance of Jaxadi is described in the [benchmarks](benchmarks/README.md) directory.
8896

89-
<!-- ## Contributing
97+
## Contributing
9098

91-
We welcome contributions! Please see our [Contributing Guide](CONTRIBUTING.md) for more details. -->
99+
We welcome contributions! Please see our [Contributing Guide](CONTRIBUTING.md) for more details.
100+
101+
## Citation
102+
103+
If you use JaxADi in your research, please cite it as follows:
104+
105+
```bibtex
106+
@misc{jaxadi2024,
107+
title = {JaxADi: Bridging CasADi and JAX for Efficient Numerical Computing},
108+
author = {Alentev, Igor and Kozlov, Lev and Nedelchev, Simeon},
109+
year = {2024},
110+
url = {https://github.com/based-robotics/jaxadi},
111+
note = {Accessed: [Insert Access Date]}
112+
}
113+
```
92114

93115
## Acknowledgements
94116

@@ -98,5 +120,4 @@ This project draws inspiration from [cusadi](https://github.com/se-hwan/cusadi),
98120

99121
For questions, issues, or suggestions, please [open an issue](https://github.com/based-robotics/jaxadi/issues) on our GitHub repository.
100122

101-
102123
We hope JAXADI empowers your numerical computing and optimization tasks! Happy coding!

benchmarks/README.md

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,31 @@
1+
# Benchmarking
2+
3+
In order to evaluate the performance of the `jaxadi` library vs `cusadi` we have tried to reproduce the benchmarks from the `cusadi` library first and faced some issues with the proper `cuda` installation.
4+
5+
![meme](https://preview.redd.it/explain-please-v0-ma2mz5wxftod1.jpeg?auto=webp&s=2b90dfa3b12e064f54333e1080b3dabbad914f48)
6+
7+
Adding the complexity of setup of benchmarks of [cusadi](https://github.com/se-hwan/cusadi) we have copied and modified the benchmarks to be able to run them in the [`colab` environment](https://colab.research.google.com/github/based-robotics/jaxadi/blob/feature%2Fbenchmarking/benchmarks/jaxadi_vs_cusadi.ipynb) side by side (CUDA vs Jax).
8+
9+
Due limitations we cover only the functions with less than 1e3 operations. All of them are located in the [data](data) directory.
10+
11+
One may run the benchmarks in the colab environment and get the [cusadi results](cuda_benchmark_results.npz) and [jaxadi results](jax_benchmark_results.npz) for comparison.
12+
13+
#### Overview
14+
15+
This analysis compares the performance of two computational tools, **Jaxadi** and **Cusadi**, using a range of batch sizes for different operations. The tools are evaluated based on average computation times and their relative speedup across various workloads.
16+
17+
![1e1](https://github.com/based-robotics/jaxadi/blob/master/docs/static/images/compare_1e1_bar.png?raw=true)
18+
19+
![1e2](https://github.com/based-robotics/jaxadi/blob/master/docs/static/images/compare_1e2_bar.png?raw=true)
20+
21+
![Speedup](https://github.com/based-robotics/jaxadi/blob/master/docs/static/images/speedup_ratio.png?raw=true)
22+
23+
#### Key Findings
24+
25+
- **Jaxadi's Consistent Advantage:** Across all tested scenarios, Jaxadi consistently shows better performance than Cusadi. This is evident from the lower average computation times for Jaxadi, particularly as the workload increases.
26+
27+
- **Scalability Differences:** As batch sizes grow, the gap between the two tools becomes more pronounced. Cusadi's computation time increases rapidly with larger batch sizes, especially beyond \(2^{12}\). In contrast, Jaxadi maintains a more stable performance with a much slower increase in computation time, indicating better batch scalability.
28+
29+
#### Conclusion
30+
31+
Overall, Jaxadi offers a clear performance advantage over Cusadi, particularly for small functions which can be translated propely.

benchmarks/cuda_benchmark_results.npz

7.71 KB
Binary file not shown.

benchmarks/data/fn_1e1.casadi

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
jhpnnagiieahaaaadaaaaaaaaaaaaaaaaafaegaakaaaaaaadfifgefhogdgehjgpgogcaaaaaaagaaaaaaaggogpfbdfgbdaaaabaaabagaaaaaaaeaaaaaaaaaaaaaaababababadaaaaaaaaaaaaaaabababaeaaaaaaaaaaaaaaaegiaaaaaaaaaaaaaaaeaaaaaaaaaaaaaaabaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaeaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaabaaaaaaaaaaaaaaacaaaaaaaaaaaaaaadaaaaaaaaaaaaaaaeggaaaaaaaaaaaaaaacaaaaaaaaaaaaaaabaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaacaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaabaaaaaaaaaaaaaaachbaaaaaaaaaaaaaaaegfaaaaaaaaaaaaaaabaaaaaaaaaaaaaaabaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaabaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaadaaaaaaaaaaaaaaachbaaaaaaaaaaaaaaaegpaaaaaaaaaaaaaaadaaaaaaaaaaaaaaadaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaadaaaaaaaaaaaaaaagaaaaaaaaaaaaaaajaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaabaaaaaaaaaaaaaaacaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaabaaaaaaaaaaaaaaacaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaabaaaaaaaaaaaaaaacaaaaaaaaaaaaaaaeghaaaaaaaaaaaaaaadaaaaaaaaaaaaaaabaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaadaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaabaaaaaaaaaaaaaaacaaaaaaaaaaaaaaaeaaaaaaaaaaaaaaacaaaaaaajgadcaaaaaaajgbdcaaaaaaajgcdcaaaaaaajgdddaaaaaaaaaaaaaaacaaaaaaapgadcaaaaaaapgbdcaaaaaaapgcdaabagaaaaaaadhpgfhchdgfgbahaaaaaaakgjgehpfehngahaaaaaaaaaaaaaaaafaaaaaaadhigfgmgmgaaaaaaaaaaaaaaaaaaegbaaaaaaaaaaaaaaaaebababaaabababaaapbfilobfilobfnpdmfpicmfpicmfpnpdaaaaaeaaaaaaaaaaaaaabakdmiadcooijhfeodaaaaaaaaaaaaabhcaaaaaaaaaaaaaaaabaaaaaaaocdaaaaaaangehihaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaachfaaaaaaaaaaaaaaaeaaaaaaaaaaaaaaadaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaacbaaaaaaaaaaaaaabaaaaaaaeaaaaaaaaaaaaaaachaaaaaaaaaaaaaaaaeaaaaaaaaaaaaaaaegpcaaaaaaaaaaaaaadaaaaaaaiepfadegpcaaaaaaaaaaaaaadaaaaaaaiepfbdegpcaaaaaaaaaaaaaadaaaaaaaiepfcdegpcaaaaaaaaaaaaaadaaaaaaaiepfddchbaaaaaaaaaaaaaaacaaaaaaaaaaaaaaaegpcaaaaaaaaaaaaaadaaaaaaadgpfadegpcaaaaaaaaaaaaaadaaaaaaadgpfbdchbaaaaaaaaaaaaaaacaaaaaaaaaaaaaaaegpcaaaaaaaaaaaaaadaaaaaaabepfadegpcaaaaaaaaaaaaaadaaaaaaabepfbdchcaaaaaaaaaaaaaaabaaaaaaaaaaaaaaaegpcaaaaaaaaaaaaaabaaaaaaacgbaaaaaaahdaaaaaaaaaaaaaacbaaaaaaaaaaaaaaaaaaaaaaaaaaaaaapbaaaaaaaaaaaaaaegbaaaaaaaaaaaaaaachgaaaaaaaaaaaaaaaegmcaaaaaaaaaaaaaachllnlhnjnpnmhlnndegeaaaaaaaaaaaaaaachkaaaaaaaaaaaaaaachabaaaaaaaaaaaaaaegeaaaaaaaaaaaaaaachmaaaaaaaaaaaaaaachabaaaaaaaaaaaaaaegdaaaaaaaaaaaaaaachcbaaaaaaaaaaaaaachkaaaaaaaaaaaaaaaegbaaaaaaaaaaaaaaachoaaaaaaaaaaaaaaachdbaaaaaaaaaaaaaaegeaaaaaaaaaaaaaaachiaaaaaaaaaaaaaaachabaaaaaaaaaaaaaaegdaaaaaaaaaaaaaaachfbaaaaaaaaaaaaaachmaaaaaaaaaaaaaaaegcaaaaaaaaaaaaaaachnaaaaaaaaaaaaaaachgbaaaaaaaaaaaaaaegbaaaaaaaaaaaaaaachjaaaaaaaaaaaaaaachpaaaaaaaaaaaaaaaegdaaaaaaaaaaaaaaachiaaaaaaaaaaaaaaachfbaaaaaaaaaaaaaaegcaaaaaaaaaaaaaaachibaaaaaaaaaaaaaachjbaaaaaaaaaaaaaaegeaaaaaaaaaaaaaaachhbaaaaaaaaaaaaaachkbaaaaaaaaaaaaaaegdaaaaaaaaaaaaaaachfbaaaaaaaaaaaaaachkaaaaaaaaaaaaaaaegcaaaaaaaaaaaaaaachmbaaaaaaaaaaaaaachlaaaaaaaaaaaaaaaegdaaaaaaaaaaaaaaachlbaaaaaaaaaaaaaachnbaaaaaaaaaaaaaaegcaaaaaaaaaaaaaaachebaaaaaaaaaaaaaachobaaaaaaaaaaaaaaegdaaaaaaaaaaaaaaachmaaaaaaaaaaaaaaachcbaaaaaaaaaaaaaaegcaaaaaaaaaaaaaaachpaaaaaaaaaaaaaaachacaaaaaaaaaaaaaaegdaaaaaaaaaaaaaaachhbaaaaaaaaaaaaaachlbaaaaaaaaaaaaaaegcaaaaaaaaaaaaaaachbcaaaaaaaaaaaaaachccaaaaaaaaaaaaaaegeaaaaaaaaaaaaaaachpbaaaaaaaaaaaaaachdcaaaaaaaaaaaaaaegdaaaaaaaaaaaaaaachcbaaaaaaaaaaaaaachecaaaaaaaaaaaaaaegbaaaaaaaaaaaaaaachbbaaaaaaaaaaaaaachfcaaaaaaaaaaaaaaegeaaaaaaaaaaaaaaachnbaaaaaaaaaaaaaachkbaaaaaaaaaaaaaaegdaaaaaaaaaaaaaaachlbaaaaaaaaaaaaaachecaaaaaaaaaaaaaaegcaaaaaaaaaaaaaaachhcaaaaaaaaaaaaaachicaaaaaaaaaaaaaaegdaaaaaaaaaaaaaaachfbaaaaaaaaaaaaaachjcaaaaaaaaaaaaaaegbaaaaaaaaaaaaaaachgcaaaaaaaaaaaaaachkcaaaaaaaaaaaaaaegfaaaaaaaaaaaaaaachlcaaaaaaaaaaaaaaegfaaaaaaaaaaaaaaachkaaaaaaaaaaaaaaaegfaaaaaaaaaaaaaaachlaaaaaaaaaaaaaaabaaaaaaaaaaaaaaachpaaaaaaaaaaaaaaaeaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaancaaaaaaaaaaaaaabaaaaaaaaaaaaaaancaaaaaabaaaaaaaaaaaaaaaaaaaaaaamcaaaaaacaaaaaaallnlhnjnpnmhlnndbaaaaaaabaaaaaaabaaaaaaacaaaaaaaeaaaaaaadaaaaaaaaaaaaaaabaaaaaaancaaaaaaeaaaaaaacaaaaaaaaaaaaaaaeaaaaaaafaaaaaaaeaaaaaaabaaaaaaancaaaaaagaaaaaaadaaaaaaaaaaaaaaadaaaaaaahaaaaaaafaaaaaaaaaaaaaaabaaaaaaahaaaaaaagaaaaaaahaaaaaaancaaaaaaiaaaaaaacaaaaaaabaaaaaaancaaaaaajaaaaaaaaaaaaaaacaaaaaaaeaaaaaaakaaaaaaajaaaaaaabaaaaaaadaaaaaaalaaaaaaakaaaaaaaeaaaaaaacaaaaaaalaaaaaaaiaaaaaaalaaaaaaancaaaaaamaaaaaaaaaaaaaaadaaaaaaabaaaaaaamaaaaaaamaaaaaaacaaaaaaadaaaaaaanaaaaaaajaaaaaaakaaaaaaacaaaaaaanaaaaaaamaaaaaaanaaaaaaaeaaaaaaaoaaaaaaalaaaaaaanaaaaaaadaaaaaaapaaaaaaakaaaaaaaaaaaaaaancaaaaaaabaaaaaabaaaaaaabaaaaaaacaaaaaaapaaaaaaapaaaaaaaabaaaaaadaaaaaaabbaaaaaaoaaaaaaapaaaaaaacaaaaaaahaaaaaaahaaaaaaabbaaaaaadaaaaaaabbaaaaaaeaaaaaaafaaaaaaacaaaaaaabbaaaaaacaaaaaaabbaaaaaadaaaaaaalaaaaaaalaaaaaaaoaaaaaaacaaaaaaabbaaaaaabbaaaaaalaaaaaaaeaaaaaaahaaaaaaahaaaaaaabbaaaaaadaaaaaaafaaaaaaafaaaaaaahaaaaaaabaaaaaaadaaaaaaadaaaaaaafaaaaaaaeaaaaaaapaaaaaaapaaaaaaanaaaaaaadaaaaaaaoaaaaaaaoaaaaaaahaaaaaaacaaaaaaapaaaaaaapaaaaaaaoaaaaaaadaaaaaaakaaaaaaakaaaaaaapaaaaaaabaaaaaaadaaaaaaadaaaaaaakaaaaaaafaaaaaaadaaaaaaadaaaaaaadaaaaaaaocaaaaaaaaaaaaaadaaaaaaaaaaaaaaaocaaaaaaaaaaaaaapaaaaaaabaaaaaaaocaaaaaabaaaaaaabaaaaaaaaaaaaaaancaaaaaabaaaaaaaaaaaaaaabaaaaaaaocaaaaaabaaaaaaabaaaaaaabaaaaaaaocaaaaaabaaaaaaaeaaaaaaacaaaaaaaocaaaaaabaaaaaaajaaaaaaadaaaaaaaocaaaaaabaaaaaaamaaaaaaaeaaaaaaaocaaaaaabaaaaaaaiaaaaaaafaaaaaaaocaaaaaabaaaaaaaeaaaaaaagaaaaaaaocaaaaaabaaaaaaaiaaaaaaahaaaaaaaocaaaaaabaaaaaaacaaaaaaaiaaaaaaafaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaocaaaaaacaaaaaaaaaaaaaaaaaaaaaaafaaaaaaaabaaaaaaabaaaaaaabaaaaaaocaaaaaacaaaaaaaabaaaaaabaaaaaaaocaaaaaacaaaaaaagaaaaaaacaaaaaaabadaaaaaaaaaaaaaaachbaaaaaaaaaaaaaaacaaaaaaaaaaaaaaachmcaaaaaaaaaaaaaachjcaaaaaaaaaaaaaachdaaaaaaaaaaaaaaajaaaaaaaaaaaaaaachabaaaaaaaaaaaaaachhaaaaaaaaaaaaaaachmaaaaaaaaaaaaaaachiaaaaaaaaaaaaaaachibaaaaaaaaaaaaaachnaaaaaaaaaaaaaaachmaaaaaaaaaaaaaaachnaaaaaaaaaaaaaaachpaaaaaaaaaaaaaaacheaaaaaaaaaaaaaaadaaaaaaaaaaaaaaachncaaaaaaaaaaaaaachocaaaaaaaaaaaaaachoaaaaaaaaaaaaaaa

0 commit comments

Comments
 (0)