2 ms·
Strange that the author claims Jax's vmap is what's doing the heavy lifting, but doesn't use PyTorch vmap to make the benchmark comparable.
by ipsum2 4y ago
Strange that the author claims Jax's vmap is what's doing the heavy lifting, but doesn't use PyTorch vmap to make the benchmark comparable.