import torch


def batch_fn(fn):
    """Transforms a function :fn: to act on each element of a batch individually."""
    def batched_fn(x, **kwargs):
        return torch.stack([fn(xi, **kwargs) for xi in x])
    return batched_fn