import torch


def get_spixel_prob(spixel_x, spixel_y):
    """
        generate soft assignment matrix via probability vector in x and y direction

        Args:
            spixel_x: torch.tensor
                A trainable tensor indicating the probability that pixels belong
                to superpxiels in x direction
            spixel_y: torch.tensor
                A trainable tensor indicating the probability that pixels belong
                to superpxiels in y direction

        Return:
            spixel_prob: torch.tensor
                A trainable tensor indicating the probability that pixels belong
                to superpixels
    """
    b, _, h, w = spixel_x.shape
    return torch.einsum("bxhw,byhw->bxyhw", spixel_x,
                        spixel_y).view(b, -1, h, w).contiguous()