import torch


def weighted_l2(y_item, y_pred, weights):
    """
    calculate weighted l2 loss for pytorch input

    Arguments:
    ----------
    y_item : torch array (N*,)
        array of true leave values (either probability or regression values)
    y_pred : torch array (N*,)
        array of predicted values for each leaf of the x input (either
        probability or regression values)
    weights_item : torch array (N*,)
        array of weights associated with each leaf of x input

    Returns:
    --------
    loss : torch scalar
        l2 loss
    """
    loss = torch.sum( weights * (y_item - y_pred)**2) / torch.sum( weights )
    return loss