Long term support of jaxlie?

[mjo22]

Hello! Me again. Again, this is a great library and thank you for responding to pull requests so promptly. I was wondering though its plans for long-term support, as I’m debating if I should keep this as a dependency in my library or write something home-cooked.

[brentyi]

Hi Michael, thanks again for the nice words!

There are no new features planned, but I'm planning to support jaxlie indefinitely. I'm also using both this JAX version and ports to PyTorch + NumPy¹ on an almost daily basis, so if you run into any issues I'd of course be very interested. 🙂

That said, jaxlie is fairly lightweight and building something yourself could make sense if you want to reduce dependencies, make breaking API changes, or are just looking for an exercise. In terms of the API, there are certainly design details that I've gone back and forth on, like:

• Putting parameters for Lie groups into a single vector always. There are usability advantages of the current system, but a lot of operations could be more efficient if wxyz_xyz in SE3(), for example, was instead two vectors (wxyz, translation) or seven (qw, qx, qy, qz, tx, ty, tz).
• Making each Lie group a dataclass. I could see some things like vmap usage and/or general batching within the PyTree to be confusing for people.
• Related: whether the Lie group abstraction is a good idea to begin with. It's second nature for a lot of us, but I think a convincing argument could be made for those who don't care as much about the math that it ends up introducing a lot of unnecessary jargon to common transformation operations. If you want to convert a quaternion to its axis-angle representation, for example, our current syntax is axis_angle = SO3(quat).log(). We could also just implement a function like axis_angle = compute_axis_angle_from_quaternion(quat) or Rotation3.from_quat(quat).to_axis_angle().

Overall I think the specific decisions we made for these things were correct given my personal priorities, but I can see other folks having different ones!

Footnotes

1. the PyTorch version isn't public, but there's a NumPy port that we're depending on pretty heavily for some visualization tooling: https://viser.studio/transforms/ ↩

[mjo22]

Thank you for answering this, this is all helpful to know.

• In my case, I’ve been wrapping the SO(3) object into something somewhat akin to the SE(3) object. On this note, are you no longer than thinking about implementing Consider/benchmark splitting group parameters into separate arrays #7? For my case, I may be interested in parameterizing by Euler angles and only batching over one angle. It would save some memory to do this.
• Batching within the PyTree does seem to be a learning curve within JAX, but my mindset has become that it’s part of the functional programming learning curve. I’ve built my package on equinox, which has some really nice utilities for batching.
• I hear this, the Lie group abstraction makes sense for this package I think! The objects I’ve been writing avoid the Lie group abstraction because it does not align with conventions of how people think about rotations in my field.

As for feature requests, if we go forward with this as a dependency at some point I would consider a PR that allows for utilities for using different euler angle conventions. However, this is easy enough to do with the package as is.

Thanks again!

[brentyi]

Hi sorry for the late reply, this got buried by holiday things...

Yeah, I don't have plans to implement #7. This would be a breaking API change, the ergonomics of having several floating arrays now feels worse to me, and I don't think the hypothetical + likely small runtime benefits would be worth the effort.

If you have features you want I'm happy to review PRs! Euler angles are frowned upon in most of the circles I find myself in (because of singularities, the illusion of interpretability, too many ordering conventions, etc), but I do think incorporating more support for them makes sense for this package. 😛

[mjo22]

Great, thanks for the update and no problem on the delay. The help is much appreciated!