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
Code: github.com/samuela/torc...

Install with `pip install torch2jax` to get started!
GitHub - samuela/torch2jax: Run PyTorch in JAX. 🤝
Run PyTorch in JAX. 🤝. Contribute to samuela/torch2jax development by creating an account on GitHub.
github.com
February 27, 2025 at 2:17 AM
Adding support for random ops and batch norm required a near complete rewrite of the library internals. v0.1.0 now makes extensive use of PyTorch Modes (pytorch.org/docs/stable/...), an underrating part of the PyTorch API IMHO.

Shout out to Nick Boyd who inspired me to undertake this rewrite!
Extending PyTorch — PyTorch 2.6 documentation
pytorch.org
February 27, 2025 at 2:17 AM
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