Broadcasting Compare

I have following little example with the variable bboxes of shape NxMxWxW, where N is my batchsize, M is the max number of bounding boxes per image, W is the Width and Height of the Image. A BBox is represented as the class idx in each pixel the bbox is present. Additional I have the variable candidate_classes of shape NxO, where O is the number of candidates. If I run the following:

bboxes_classes = bboxes.amax((2,3))

I will get for each bbox the class label in shape NxM. Where I now struggle is with my comparison. The goal is to get a Tensor with shape NxMxO with some vectorizing shenanigans, so that I have compared for each image and for each bbox its class against each candidates class. If for loops are necessary I don’t need any help since this is already implemented.