import torch


def compute_iou(boxes1, boxes2):
    """Compute IoU of two sets of boxes, each box is [x1, y1, x2, y2].

    Arguments:
        boxes1: a float tensor of shape [n, 4].
        boxes2: a float tensor of shape [m, 4].

    Returns:
        a float tensor of shape [n, m].
    """
    n = boxes1.size(0)
    m = boxes2.size(0)

    # left top
    lt = torch.max(
        boxes1[:, :2].unsqueeze(1).expand(n, m, 2),
        boxes2[:, :2].unsqueeze(0).expand(n, m, 2),
    )
    # [n, 2] -> [n, 1, 2] -> [n, m, 2]
    # [m, 2] -> [1, m, 2] -> [n, m, 2]

    # right bottom
    rb = torch.min(
        boxes1[:, 2:].unsqueeze(1).expand(n, m, 2),
        boxes2[:, 2:].unsqueeze(0).expand(n, m, 2),
    )

    # width height
    wh = rb - lt  # [n, m, 2]
    wh[wh < 0.0] = 0.0
    inter = wh[:, :, 0] * wh[:, :, 1]  # [n, m]

    area1 = (boxes1[:, 2] - boxes1[:, 0])*(boxes1[:, 3] - boxes1[:, 1])  # [n]
    area2 = (boxes2[:, 2] - boxes2[:, 0])*(boxes2[:, 3] - boxes2[:, 1])  # [m]
    area1 = area1.unsqueeze(1).expand_as(inter)  # [n] -> [n, 1] -> [n, m]
    area2 = area2.unsqueeze(0).expand_as(inter)  # [m] -> [1, m] -> [n, m]

    iou = inter/(area1 + area2 - inter)
    return iou