[chillee]

I'd also note that you can more or less write the same code in PyTorch with torch.vmap