from __future__ import annotations

import torch


def masked_mse(pred: torch.Tensor, target: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
    safe_diff = torch.where(mask > 0, pred - target, torch.zeros_like(pred))
    error = (safe_diff ** 2) * mask
    denom = torch.clamp(mask.sum(), min=1.0)
    return error.sum() / denom


def masked_mae(pred: torch.Tensor, target: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
    safe_diff = torch.where(mask > 0, pred - target, torch.zeros_like(pred))
    error = torch.abs(safe_diff) * mask
    denom = torch.clamp(mask.sum(), min=1.0)
    return error.sum() / denom
