import torch


def alpha_composite(base, added, eps=1e-8):
    """
    Composite two tensors, i.e., layers `added` on top of `base`,
    where the last channel is assumed to be an alpha channel.
    [... C H W], [... C H W] -> [... C H W]
    """
    # Separate color and alpha
    alpha_b =  base[..., -1:, :, :]
    alpha_a = added[..., -1:, :, :]

    color_b =  base[..., :-1, :, :]
    color_a = added[..., :-1, :, :]

    # https://en.wikipedia.org/wiki/Alpha_compositing#Alpha_blending
    alpha_0 = (1 - alpha_a) * alpha_b + alpha_a
    color_0 = ((1-alpha_a) * alpha_b*color_b + alpha_a*color_a) / (alpha_0 + eps)

    # Re-combine new color and alpha
    return torch.concat([color_0, alpha_0], dim=-3)