from typing import Callable
import torch


def apply(func: Callable, M: torch.Tensor, *args, **kwargs) -> torch.Tensor:
    """Apply a function to each `row` of a tensor. In particular, if M is a NxHx... tensor
    Applies the function func to each tensor of size Hx...
    and return a tensor of size NxOx...

    Args:
        func (callable): function that takes inputs of size Hx...
        M (torch.tensor): NxHx...tensor
        args, kwargs: passed to func
    Returns:
        A torch tensor object where the function func has been to each row
    """
    # https://discuss.pytorch.org/t/a-fast-way-to-apply-a-function-across-an-axis/8378
    tList = [func(m, *args, **kwargs) for m in torch.unbind(M, dim=0)]
    res = torch.stack(tList, dim=0)
    return res