import torch


def maximum_likelihood_loss(z, ldj):
    """
    Maximum likelihood loss
    :param z: Tensor of shape (batch_size, horizon, z_dim) -- the output of the final block
    :param ldj: Tensor of shape (batch_size, horizon, 1) -- the log determinant of the jacobian.
    :return: ML loss (scalar)
    """
    return torch.sum(
        torch.mean(
            ((0.5 * torch.norm(z, dim=-1, keepdim=True).pow(2.0)) - ldj).squeeze(-1),
            dim=-1,
        ),
        dim=-1,
    )