lucidrains / jax2torchLinks
Use Jax functions in Pytorch
β251Updated 2 years ago
Alternatives and similar repositories for jax2torch
Users that are interested in jax2torch are comparing it to the libraries listed below
Sorting:
- Run PyTorch in JAX. π€β294Updated last month
- Running Jax in PyTorch Lightningβ113Updated 9 months ago
- Pytorch-like dataloaders for JAX.β94Updated 4 months ago
- A library for programmatically generating equivariant layers through constraint solvingβ275Updated 2 years ago
- Flow-matching algorithms in JAXβ105Updated last year
- JMP is a Mixed Precision library for JAX.β208Updated 8 months ago
- OpTree: Optimized PyTree Utilitiesβ195Updated this week
- A functional training loops library for JAXβ88Updated last year
- LoRA for arbitrary JAX models and functionsβ142Updated last year
- JAX Arrays for human consumptionβ106Updated 3 months ago
- β115Updated 3 weeks ago
- β158Updated last year
- A simple library for scaling up JAX programsβ143Updated 11 months ago
- A parallel ODE solver for PyTorchβ266Updated last year
- Lightning-like training API for JAX with Flaxβ42Updated 9 months ago
- A convenient way to trigger synchronizations to wandb / Weights & Biases if your compute nodes don't have internet!β85Updated last month
- Graph neural networks in JAX.β67Updated last year
- β120Updated 3 months ago
- CLU lets you write beautiful training loops in JAX.β356Updated 3 months ago
- Simple tools to mix and match PyTorch and Jax - Get the best of both worlds!β33Updated 5 months ago
- Second Order Optimization and Curvature Estimation with K-FAC in JAX.β288Updated last week
- Flow Annealed Importance Sampling Bootstrap (FAB). ICLR 2023.β65Updated last year
- Bare-bones implementations of some generative models in Jax: diffusion, normalizing flows, consistency models, flow matching, (beta)-VAEsβ¦β133Updated last year
- Implementing the Denoising Diffusion Probabilistic Model in Flaxβ151Updated 2 years ago
- Equivariant Steerable CNNs Library for Pytorch https://quva-lab.github.io/escnn/β31Updated 2 years ago
- A Python package of computer vision models for the Equinox ecosystem.β109Updated last year
- Wraps PyTorch code in a JIT-compatible way for JAX. Supports automatically defining gradients for reverse-mode AutoDiff.β58Updated 2 months ago
- Implementations and checkpoints for ResNet, Wide ResNet, ResNeXt, ResNet-D, and ResNeSt in JAX (Flax).β112Updated 3 years ago
- NF-Layers for constructing neural functionals.β89Updated last year
- Implementation of Denoising Diffusion Probabilistic Models (DDPM) in JAX and Flax.β20Updated last year