Samuel Ainsworth
banner
skainswo.bsky.social
Samuel Ainsworth
@skainswo.bsky.social
prev: @BrownUniversity, @uwcse/@uw_wail phd, ex-@cruise, RS @waymo. 0.1x engineer, 10x friend.

spondyloarthritis, cars ruin cities, open source
As a reminder, torch2jax enables running PyTorch code in JAX. Mix-and-match PyTorch and JAX code with seamless, end-to-end autodiff, use JAX classics like jit, grad, and vmap on PyTorch code, and run PyTorch models on TPUs.

See x.com/SamuelAinswo... for the initial project announcement.
February 27, 2025 at 2:17 AM
I'm excited to announce the release of torch2jax v0.1.0!

Now with support for random functions like `torch.rand` and modules that use mutable buffers like BatchNorm 🎲
February 27, 2025 at 2:17 AM