import torch


def edge_clamp(B_cp, edge_idx, clamp_val=0, device='cpu', dtype=torch.float32):
    """
    Clamp edges of each component of Bcp to 
    clamp_val.
    RH 2020
    
    Args:
        B_cp (list of torch.Tensor):
            Beta Kruskal tensor.
            List of tensors of shape
             (n_features, rank).
    
    Returns:
        B_cp_clamped (list of torch.Tensor):
            Beta Kruskal tensor with edges clamped.
    """
    eIdxBool = torch.ones(B_cp[0].shape[0], dtype=dtype, device=device)
    eIdxBool[edge_idx] = clamp_val
    return [B_cp[0]*eIdxBool[:,None,None] , B_cp[1], B_cp[2]]