Does it need to be differentiable?
Otherwise:
a = torch.randint(0, 4, [2,3,4])
b = torch.randint(0, 4, [2,3,4])
print(a)
print(b)

print(torch.unique(a[a == b], return_counts=True))

output:
tensor([[[2, 1, 2, 1],
         [0, 0, 0, 0],
         [0, 0, 0, 1]],

        [[1, 3, 2, 3],
         [3, 3, 0, 0],
         [3, 1, 3, 3]]])

tensor([[[0, 3, 1, 3],
         [2, 0, 3, 3],
         [2, 0, 3, 1]],

        [[2, 2, 2, 1],
         [0, 2, 3, 0],
         [1, 0, 0, 2]]])

(tensor([0, 1, 2]), tensor([3, 1, 1]))