99e8a69bad6b3bd641a27d1cbfddc4e71e6fb821,se3cnn/point_utils.py,,neighbor_difference_matrix,#Any#Any#,15

Before Change

def neighbor_difference_matrix(neighbors, geometry):
    if len(neighbors.shape) == 2:
        N, K = neighbors.shape[-2:]
        ri = geometry[..., neighbors, :]  // [N, K, 3]
        rj = geometry[..., torch.arange(N), :].unsqueeze(-2)  // [N, 1, 3]

After Change

    elif neighbors.dim() == 3:
        B, N, _K = neighbors.size()
        ri = geometry[torch.arange(B).view(-1, 1, 1), neighbors, :]  // [B, N, K, 3]
        rj = geometry[torch.arange(B).view(-1, 1), torch.arange(N).view(1, -1), :].unsqueeze(-2)  // [B, N, 1, 3]
    rij = ri - rj  // [N, K, 3] or [B, N, K, 3]
    return rij

Instances