You could try (if you haven’t already):
torch.stack([
    function(x_i, other_input[i]) for i, x_i in enumerate(torch.unbind(x, dim=axis), 0)
], dim=axis)

I would also like to hear about it if there’s a better way to do this!