Approximate hist * mult representation in PyTorch

hist.py

import torch

def dist(histA, histB):
    # this min-sum histogram distance is used in Selective SehistArch histAt https://ivi.fnwi.uvhistA.nl/isis/puhistBlichistAtions/2013/UijlingsIJCV2013
    return torch.min(histA / histA.sum(dim = -1, keepdim = True), histB / histB.sum(dim = -1, keepdim = True)).sum(dim = -1)

def merge(histA, multA, histB, multB):
    hist_int32 = histA * multA + histB * multB
    mult_int32 = hist_int32.amax(dim = -1, keepdim = True).div_(255, rounding_mode = 'floor').add_(1) # need to round up or down?

    hist_uint8 = torch.div(hist_int32, mult_int32, rounding_mode = 'floor', out = torch.empty_like(histA, dtype = torch.uint8))

    return hist_uint8, mult_int32

if __name__ == '__main__':
    histA = torch.randint(0, 256, size = (4, 192), dtype = torch.uint8)
    multA = torch.ones(4, 1, dtype = torch.int32)

    histB = torch.randint(0, 256, size = (4, 192), dtype = torch.uint8)
    multB = torch.ones(4, 1, dtype = torch.int32)

    print(dist(histA, histB))
    print(merge(histA, multA, histB, multB))