def forward(self, x: QTensor) -> torch.Tensor:
"""Computes the forward pass to convert a quaternion vector to a real vector"""
if self.type == "sum":
return x.r + x.i + x.j + x.k
elif self.type == "mean":
return (x.r + x.i + x.j + x.k).mean()
elif self.type == "norm":
return x.norm()
else:
x = torch.cat([x.r, x.i, x.j, x.k], dim=-1)
return self.affine(x)
def interaction(q: QTensor) -> torch.Tensor:
norm = q.norm()
c = norm.mean(dim=-1).unsqueeze(1)
c = c.expand_as(norm)
f = norm / torch.max(norm, c)
return f
