search

ami-iit / adam

adam implements a collection of algorithms for calculating rigid-body dynamics in Jax, CasADi, PyTorch, and Numpy.
https://adam-docs.readthedocs.io/en/latest/
BSD 3-Clause "New" or "Revised" License
124 stars 19 forks source link

Consider generalizing backends of array APIs #98

Open diegoferigo opened 1 day ago

diegoferigo commented 1 day ago

I remember a f2f discussion a while ago with @Giulero about how to generalize algorithms to work on NumPy / JAX / PyTorch etc. At that time, there was no concrete solution excluding what today became Farama-Foundation/Jumpy that was only supporting NumPy and JAX.

After a couple of years, it seems that the community is trying to address this use case with the Python array API standard. It seems that this solution is the one that got major traction and attracted most of the interest towards this unification.

ADAM, today, generalizes the underlying array APIs used by the RBDAs by introducing a custom abstraction layer based on adam.core.spatial_math.{ArrayLike|ArrayLikeFactory|SpatialMath}, which has implementations in CasADi, JAX, NumPy, and PyTorch.

I think that CasADi --at least in the short term and at least officially-- won't be compatible with these Array APIs. However, providing in ADAM a new backend based on these new APIs could ease the maintenance of the other backends (or, possibly remove them), and could bring with no additional effort also support for all other array providers that are adhere with the API standard.

A good starting point could be the data-apis/array-api-compat wrapper^1, its usage seems pretty straightforward.

Giulero commented 1 day ago

Thanks a lot @diegoferigo! <3 This would be lovely indeed. I agree that finally exploiting these new Array APIs will make the library more maintainable. Regarding CasADi, we could maybe maintain the actual abstraction layer and implement a new backend for what we can unify. A sort of hybrid structure, eventually converging to a cleaner solution when we find a solution for the CasADi backend. As soon as possible I can try to draft a PR and ask you other suggestions ;)

diegoferigo commented 23 hours ago

Regarding CasADi, we could maybe maintain the actual abstraction layer and implement a new backend for what we can unify. A sort of hybrid structure, eventually converging to a cleaner solution when we find a solution for the CasADi backend.

Yep I agree! If there is interest in merging at least the remaining ones by exploiting array-api-compact, then ADAM could probably provide a custom CasADi implementation based on their array APIs. There is no need for a full implementation, likely only a subset of array operations (those that are currently abstracted) are necessary.