def compute_target_probability(target_distributions, predicted_distributions):
    """Compute predicted probability mass that reaches a non zero probability target state

    Args:
        target_distributions: [n_pred, n_node]
        predicted_distributions: [n_pred, n_node]

    Example:
        >>> import torch
        >>> compute_target_probability(torch.tensor([[.2, .8, 0.],
        ...                                          [.1, 0., .9]]),
        ...                            torch.tensor([[.6, 0., .4],
        ...                                          [.3, .1, .6]]))
        tensor([0.6000, 0.9000])

    Returns:
        [n_pred, ]
    """
    target_mask = target_distributions > 0
    return (predicted_distributions * target_mask.float()).sum(dim=1)