Spaces:
Sleeping
Sleeping
File size: 3,062 Bytes
899c526 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 |
import lietorch_backends
import torch
import torch.nn.functional as F
class GroupOp(torch.autograd.Function):
""" group operation base class """
@classmethod
def forward(cls, ctx, group_id, *inputs):
ctx.group_id = group_id
ctx.save_for_backward(*inputs)
out = cls.forward_op(ctx.group_id, *inputs)
return out
@classmethod
def backward(cls, ctx, grad):
error_str = "Backward operation not implemented for {}".format(cls)
assert cls.backward_op is not None, error_str
inputs = ctx.saved_tensors
grad = grad.contiguous()
grad_inputs = cls.backward_op(ctx.group_id, grad, *inputs)
return (None, ) + tuple(grad_inputs)
class Exp(GroupOp):
""" exponential map """
forward_op, backward_op = lietorch_backends.expm, lietorch_backends.expm_backward
class Log(GroupOp):
""" logarithm map """
forward_op, backward_op = lietorch_backends.logm, lietorch_backends.logm_backward
class Inv(GroupOp):
""" group inverse """
forward_op, backward_op = lietorch_backends.inv, lietorch_backends.inv_backward
class Mul(GroupOp):
""" group multiplication """
forward_op, backward_op = lietorch_backends.mul, lietorch_backends.mul_backward
class Adj(GroupOp):
""" adjoint operator """
forward_op, backward_op = lietorch_backends.adj, lietorch_backends.adj_backward
class AdjT(GroupOp):
""" adjoint operator """
forward_op, backward_op = lietorch_backends.adjT, lietorch_backends.adjT_backward
class Act3(GroupOp):
""" action on point """
forward_op, backward_op = lietorch_backends.act, lietorch_backends.act_backward
class Act4(GroupOp):
""" action on point """
forward_op, backward_op = lietorch_backends.act4, lietorch_backends.act4_backward
class Jinv(GroupOp):
""" adjoint operator """
forward_op, backward_op = lietorch_backends.Jinv, None
class ToMatrix(GroupOp):
""" convert to matrix representation """
forward_op, backward_op = lietorch_backends.as_matrix, None
### conversion operations to/from Euclidean embeddings ###
class FromVec(torch.autograd.Function):
""" convert vector into group object """
@classmethod
def forward(cls, ctx, group_id, *inputs):
ctx.group_id = group_id
ctx.save_for_backward(*inputs)
return inputs[0]
@classmethod
def backward(cls, ctx, grad):
inputs = ctx.saved_tensors
J = lietorch_backends.projector(ctx.group_id, *inputs)
return None, torch.matmul(grad.unsqueeze(-2), torch.linalg.pinv(J)).squeeze(-2)
class ToVec(torch.autograd.Function):
""" convert group object to vector """
@classmethod
def forward(cls, ctx, group_id, *inputs):
ctx.group_id = group_id
ctx.save_for_backward(*inputs)
return inputs[0]
@classmethod
def backward(cls, ctx, grad):
inputs = ctx.saved_tensors
J = lietorch_backends.projector(ctx.group_id, *inputs)
return None, torch.matmul(grad.unsqueeze(-2), J).squeeze(-2)
|