Explore solutions for enhancing OOP in high-level resources

[diegoferigo]

Although JAX promotes a functional programming pattern, it's often useful offering to downstream users also OOP resources that usually simplify the learning curve.

Jaxsim follows a similar approach:

• All resources inside jaxsim.physics are functional.
• All resources inside jaxsim.high_level are object-oriented.

We are currently trying to simplify the application of jit/vmap transformations (see #44 and #48), but there are some use-cases related to r/w classes like high_level.link.Link and high_level.joint.Joint that are still a bit rough (see 48#issuecomment-1746635121).

It would be interesting exploring solutions like those discussed in google/jax#7919 and google/jax#17341, with some extra detail provided in #48 (comment).

[diegoferigo]

We eventually decided to migrate towards functional APIs. The OOP resources have been removed and merged into main in #108.

[diegoferigo]

For those reading this issue in mid 2024, there's an experimental attempt of using PyDags for OOP in Flax:

https://flax.readthedocs.io/en/v0.8.3/experimental/nnx/nnx_basics.html

In their case, the new classes hold both the NN structure and the state (parameters). A caveat -at first look- is that they had to create a new decorator (rings a bell with our previous OOP pattern) for JIT compilation to handle stateful arguments.