Comment by chillee

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