import torch


def apply_examplewise(fn, x: torch.Tensor, sequence_lengths, time_axis=1):
    """
    Applies a function to each element of x (along batch (0) dimension) and
    respects the sequence lengths along time axis. Assumes that fn does not
    change the dimensions of its input (e.g., norm).
    """
    if sequence_lengths is None:
        return fn(x)
    else:
        # Check inputs
        assert time_axis != 0, 'The first axis must be the batch axis!'
        assert len(sequence_lengths) == x.shape[0], (
            'Number of sequence lengths and batch size must match!'
        )

        time_axis = time_axis % x.dim()
        selector = [slice(None)] * (time_axis - 1)
        out = torch.zeros_like(x)
        for b, l in enumerate(sequence_lengths):
            s = (b, *selector, slice(l))

            # Keep the batch dimension while processing
            out[s] = fn(x[s][None, ...])[0]
        return out