Refactor algorithms to use array-api-compat and native batch operations

README.md

@@ -18,19 +18,22 @@
 
 ### Table of contents
 
-- [🐍 Dependencies](#-dependencies)
-- [💾 Installation](#-installation)
-  - [🐍 Installation with pip](#-installation-with-pip)
-  - [📦 Installation with conda](#-installation-with-conda)
-    - [Installation from conda-forge package](#installation-from-conda-forge-package)
-  - [🔨 Installation from repo](#-installation-from-repo)
-- [🚀 Usage](#-usage)
-  - [Jax interface](#jax-interface)
-  - [CasADi interface](#casadi-interface)
-  - [PyTorch interface](#pytorch-interface)
-  - [PyTorch Batched interface](#pytorch-batched-interface)
-- [🦸‍♂️ Contributing](#️-contributing)
-- [Todo](#todo)
+- [adam](#adam)
+    - [Table of contents](#table-of-contents)
+  - [🐍 Dependencies](#-dependencies)
+  - [💾 Installation](#-installation)
+    - [🐍 Installation with pip](#-installation-with-pip)
+    - [📦 Installation with conda](#-installation-with-conda)
+      - [Installation from conda-forge package](#installation-from-conda-forge-package)
+    - [🔨 Installation from repo](#-installation-from-repo)
+  - [🚀 Usage](#-usage)
+    - [Jax interface](#jax-interface)
+    - [CasADi interface](#casadi-interface)
+    - [PyTorch interface](#pytorch-interface)
+    - [PyTorch Batched interface](#pytorch-batched-interface)
+    - [Inverse Kinematics](#inverse-kinematics)
+  - [🦸‍♂️ Contributing](#️-contributing)
+  - [Todo](#todo)
 
 ## 🐍 Dependencies
 
@@ -43,7 +46,7 @@ Other requisites are:
 - `casadi`
 - `pytorch`
 - `numpy`
-- `jax2torch`
+- `array-api-compat`
 
 They will be installed in the installation step!
 
@@ -145,25 +148,25 @@ Install in a conda environment the required dependencies:
 - **Jax** interface dependencies:
 
   ```bash
-  conda create -n adamenv -c conda-forge jax numpy lxml prettytable matplotlib urdfdom-py
+  conda create -n adamenv -c conda-forge jax numpy lxml prettytable matplotlib urdfdom-py array-api-compat
   ```
 
 - **CasADi** interface dependencies:
 
   ```bash
-  conda create -n adamenv -c conda-forge casadi numpy lxml prettytable matplotlib urdfdom-py
+  conda create -n adamenv -c conda-forge casadi numpy lxml prettytable matplotlib urdfdom-py array-api-compat
   ```
 
 - **PyTorch** interface dependencies:
 
   ```bash
-  conda create -n adamenv -c conda-forge pytorch numpy lxml prettytable matplotlib urdfdom-py jax2torch
+  conda create -n adamenv -c conda-forge pytorch numpy lxml prettytable matplotlib urdfdom-py array-api-compat
   ```
 
 - **ALL** interfaces dependencies:
 
   ```bash
-  conda create -n adamenv -c conda-forge jax casadi pytorch numpy lxml prettytable matplotlib urdfdom-py jax2torch
+  conda create -n adamenv -c conda-forge jax casadi pytorch numpy lxml prettytable matplotlib urdfdom-py array-api-compat
   ```
 
 Activate the environment, clone the repo and install the library:
@@ -316,8 +319,6 @@ print(M)
 
 ### PyTorch Batched interface
 
-> [!NOTE]
-> When using this interface, note that the first call of the jitted function can be slow, since JAX needs to compile the function. Then it will be faster!
 
 ```python
 import adam

ci_env.yml

@@ -18,6 +18,6 @@ dependencies:
   - idyntree >=11.0.0
   - jax<=0.6.0
   - pytorch
-  - jax2torch
   - requests
+  - array-api-compat
   - liecasadi

ci_env_win.yml

@@ -18,3 +18,4 @@ dependencies:
   - idyntree >=11.0.0
   - requests
   - liecasadi
+  - array-api-compat

docs/index.rst

@@ -20,7 +20,6 @@ adam
 
 
 **adam** employs the automatic differentiation capabilities of these frameworks to compute, if needed, gradients, Jacobian, Hessians of rigid-body dynamics quantities. This approach enables the design of optimal control and reinforcement learning strategies in robotics.
-Thanks to the `jax.vmap`-ing and `jax.jit`-ing capabilities, the algorithms can be run on batches of inputs, which are possibly converted to PyTorch using the `jax2torch` conversion functions.
 
 
 **adam** is based on **Roy Featherstone's Rigid Body Dynamics Algorithms**.

docs/modules/pytorch_batched.rst

@@ -1,18 +1,7 @@
 PyTorch Batched interface
 =========================
 
-This module implements the batched version of the Rigid Body Dynamics algorithms using PyTorch.
-This module uses Jax under the hood, then the functions are ``jax.vmap``-ed and ``jax.jit``-ed to run on batches of inputs, which are ultimately converted to PyTorch using the ``jax2torch`` conversion functions.
-
-.. note::
-
-   The first time you run a function from this module, it will take a bit longer to execute as they are being compiled by Jax.
-
-
-.. note::
-
-   If the GPU support for ``JAX`` is needed, follow the instructions in the `Jax documentation <https://jax.readthedocs.io/en/latest/installation.html#conda-community-supported>`_.
-
+This module implements the batched version of the Rigid Body Dynamics algorithms using PyTorch. 
 
 ----------------
 

pyproject.toml

@@ -31,12 +31,13 @@ dependencies = [
   "casadi>=3.6",
   "prettytable",
   "urdfdom-py",
+  "array-api-compat",
 ]
 
 [project.optional-dependencies]
 jax = ["jax<=0.6.0", "jaxlib"]
 casadi = ["casadi"]
-pytorch = ["torch", "jax", "jaxlib", "jax2torch"]
+pytorch = ["torch", "jax", "jaxlib"]
 test = [
   "jax<=0.6.0",
   "jaxlib",
@@ -46,12 +47,11 @@ test = [
   "idyntree",
   "icub-models",
   "black",
-  "jax2torch",
   "requests",
   "liecasadi",
 ]
 conversion = ["idyntree"]
-all = ["jax", "jaxlib", "casadi>=3.6", "torch", "jax2torch"]
+all = ["jax", "jaxlib", "casadi>=3.6", "torch"]
 
 [project.urls]
 "Documentation" = "https://adam-robotics.readthedocs.io/en/latest/"