jac_large_angle = phi.__class__(
len(large_angle_inds), cls.dim, cls.dim)
jac_large_angle[:, 0, 0] = A
jac_large_angle[:, 0, 1] = -B
jac_large_angle[:, 1, 0] = B
jac_large_angle[:, 1, 1] = A
jac[large_angle_inds] = jac_large_angle
After Change
if len(large_angle_inds) > 0:
angle = phi[large_angle_inds]
s = angle.sin()
c = angle.cos()
A = (s / angle).unsqueeze_(dim=1).unsqueeze_(
dim=2).expand_as(jac[large_angle_inds]) * \
torch.eye(cls.dim).unsqueeze_(dim=0).expand_as(
jac[large_angle_inds])
B = ((1. - c) / angle).unsqueeze_(dim=1).unsqueeze_(
dim=2).expand_as(jac[large_angle_inds]) * \
cls.wedge(phi.__class__([1.]))
jac[large_angle_inds] = A + B
return jac.squeeze_()
@classmethod