Make LieTensor a subclass of torch.Tensor and override function/operators

[luisenp]

This PR is mostly to illustrate some ideas on how to safely overload operators in LieTensor class. It'd better to start looking at the example script, and if that looks OK, you can take a quick look at what's happening under the hood. This is not meant to be a final implementation.

[mhmukadam]

The example script looks good to me!

[luisenp]

One other detail to consider for discussion. I did some quick tests with torch optimizers, and could only get the subclass based version to work by doing the following, which requires an extra step. We also need to do some tests to make sure backprop is doing the right thing. Curious to hear thoughts from @albanD

import theseus.labs.lie as lie
import theseus.labs.lie.functional as lieF

# Leaf tensor needs to be a regular tensor, so we need to explicitly pass the tensor data
g1_data = lieF.se3.rand(1, requires_grad=True)
g1 = lie.cast(g1_data, ltype=lie.SE3)
g2 = lie.rand(1, lie.SE3)
opt = torch.optim.Adam([g1_data])

opt.zero_grad()
d = g1.inv().compose(g2).log()
loss = torch.sum(d**2)
loss.backward()
opt.step()

[luisenp]

cc @rmurai0610

[fantaosha]

Some comments. Not sure if we should override __add() and use LieAsEuclidean.

[fantaosha]

Honestly, I'm a little nervous about overriding __add() since they are very dangerous. Also, one can not use *, -, /, either in that case.

[fantaosha]

Do we have some operators called to_euclidean() that makes a pytorch tensor copy (not reference) of current data.

[fantaosha]

In summary, I would like exactly function names rather than overrides to avoid confusion. I encountered lots of problems to generalize the first-order optimizers on Lie groups.

[fantaosha]

To be consistent, if g + x = g.compose(x), then g-x should be g.compose(x.inverse()).

[fantaosha]

I'm concerned with this. If the user wants to use Lie group as a Euclidean tensor. Maybe it is safer to do something like g.to_euclidean().matmul(torch.rand(4,7). Maybe to_euclidean() should be a property such that users can not modify the data.

[luisenp]

Is the concern that users would override the internal data? (e.g., something like g.add_(something)? This seems a reasonable concern, because I believe subclassing allows this. I'm OK with replacing this with to_euclidean().

[luisenp]

A less esoteric but perhaps less safe alternative for add_ could be to check the tensor shape and if it looks like (...,3, 4) (for SE3) this is interpreted as an euclidean gradient, something with shape (..., 6) (for SE3) interpret this as a Riemannian graadient, and everything else throws an error.

[mhmukadam]

Looks great! Let's merge once the open comments are resolved.

[mhmukadam]

LGTM!