While many tools use symbolical representation provided by CasADi, most of them are difficult to parallelize and compute efficiently. In the recent years high-preformance numerical computing is becoming more popular thanks to frameworks like JAX, CUDA and others.
With introduction of CasADi support in the tools we use and love, we felt rising importance of proper parallelizing of those tools. This way we developed and made publically available translation utility, which allows to convert user defined CasADi functions to JIT compilable JAX functions with ability to export them or declare immidiately.
This framework allows us to achieve result comparable with CUDA native compilations, while delegating environment and hardware support from the user to the library, which is JAX in our case. JAX JIT compilation allows to run the code on various accelerator devices without the need of user intervention.
Note: For now translation does not support functions with very large number of operations, due to the translation implementation. Secret component of translation is work-tree expansion, which might lead to large overhead in number of symbols. We are working on finding the compromise in both speed and extensive functions support. Our initial goal was combining this approach with pinocchio and CasADi in robotics applications and there was no need for functions with large dimensions. However, some work on this issue is planned.
@misc{jaxadi2024,
title = {JaxADi: Bridging CasADi and JAX for Efficient Numerical Computing},
author = {Alentev, Igor and Kozlov, Lev and Nedelchev, Simeon},
year = {2024},
url = {https://github.com/based-robotics/jaxadi},
note = {Accessed: [Insert Access Date]}
}