import torch


def autograd_individual_gradients(losses, params, concat=False):
    """Compute individual gradients from individual losses via ``torch.autograd``.

    Args:
        losses (torch.Tensor): Individual losses.
        params ([torch.Tensor]): List of torch.Tensors holding the network's
            parameters.
        concat (bool): If ``True``, flatten and concatenate the gradients over all
            parameters.

    Returns:
        Tensor: Individual gradients.
    """
    # nesting is [sample, param]
    individual_gradients = [
        torch.autograd.grad(loss, params, create_graph=True) for loss in losses
    ]

    num_layers = len(params)
    batch_axis = 0

    # concatenate samples for each parameter, nesting is [param, sample]
    individual_gradients = [
        torch.cat(
            [g[layer_idx].unsqueeze(batch_axis) for g in individual_gradients],
            dim=batch_axis,
        )
        for layer_idx in range(num_layers)
    ]

    # concatenate over parameters
    if concat:
        individual_gradients = torch.cat(
            [g.flatten(start_dim=1) for g in individual_gradients], dim=1
        )

    return individual_gradients