Use Jax functions in Pytorch with DLPack

Last update: Dec 17, 2022

Overview

jax2torch

Use Jax functions in Pytorch with DLPack, as outlined in a gist by @mattjj. The repository was made for the purposes of making this differentiable alignment work interoperable with Pytorch projects.

Install

$ pip install jax2torch

Usage

Quick test

import jax
import torch
from jax2torch import jax2torch

# Jax function

@jax.jit
def jax_pow(x, y = 2):
  return x ** y

# convert to Torch function

torch_pow = jax2torch(jax_pow)

# run it on Torch data!

x = torch.tensor([1., 2., 3.])
y = torch_pow(x, y = 3)
print(y)  # tensor([1., 8., 27.])

# And differentiate!

x = torch.tensor([2., 3.], requires_grad = True)
y = torch.sum(torch_pow(x, y = 3))
y.backward()
print(x.grad) # tensor([12., 27.])

You might also like...

A PyTorch implementation of EfficientNet

EfficientNet PyTorch Quickstart Install with pip install efficientnet_pytorch and load a pretrained EfficientNet with: from efficientnet_pytorch impor

7.2k Jan 6, 2023

A collection of extensions and data-loaders for few-shot learning & meta-learning in PyTorch

Torchmeta A collection of extensions and data-loaders for few-shot learning & meta-learning in PyTorch. Torchmeta contains popular meta-learning bench

1.7k Jan 6, 2023

PyTorch Extension Library of Optimized Scatter Operations

PyTorch Scatter Documentation This package consists of a small extension library of highly optimized sparse update (scatter and segment) operations fo

1.2k Jan 7, 2023

PyTorch Extension Library of Optimized Autograd Sparse Matrix Operations

PyTorch Sparse This package consists of a small extension library of optimized sparse matrix operations with autograd support. This package currently

757 Jan 4, 2023

Reformer, the efficient Transformer, in Pytorch

Reformer, the Efficient Transformer, in Pytorch This is a Pytorch implementation of Reformer https://openreview.net/pdf?id=rkgNKkHtvB It includes LSH

1.8k Jan 6, 2023

higher is a pytorch library allowing users to obtain higher order gradients over losses spanning training loops rather than individual training steps.

higher is a library providing support for higher-order optimization, e.g. through unrolled first-order optimization loops, of "meta" aspects of these

1.5k Jan 3, 2023

PyTorch implementation of TabNet paper : https://arxiv.org/pdf/1908.07442.pdf

README TabNet : Attentive Interpretable Tabular Learning This is a pyTorch implementation of Tabnet (Arik, S. O., & Pfister, T. (2019). TabNet: Attent

2k Dec 27, 2022

PyTorch extensions for fast R&D prototyping and Kaggle farming

Pytorch-toolbelt A pytorch-toolbelt is a Python library with a set of bells and whistles for PyTorch for fast R&D prototyping and Kaggle farming: What

1.3k Jan 5, 2023

An implementation of Performer, a linear attention-based transformer, in Pytorch

Performer - Pytorch An implementation of Performer, a linear attention-based transformer variant with a Fast Attention Via positive Orthogonal Random

900 Dec 22, 2022

Releases(0.0.7)

0.0.7(Oct 26, 2021)

Source code(tar.gz)
Source code(zip)
0.0.6(Oct 26, 2021)

Source code(tar.gz)
Source code(zip)
0.0.5(Oct 26, 2021)

Source code(tar.gz)
Source code(zip)
0.0.4(Oct 26, 2021)

Source code(tar.gz)
Source code(zip)
0.0.3(Oct 26, 2021)

Source code(tar.gz)
Source code(zip)
0.0.2(Oct 26, 2021)

Source code(tar.gz)
Source code(zip)
0.0.1a(Oct 26, 2021)

Source code(tar.gz)
Source code(zip)

Owner

Phil Wang

Working with Attention. It's all we need

GitHub Repository

High-level batteries-included neural network training library for Pytorch

Pywick High-Level Training framework for Pytorch Pywick is a high-level Pytorch training framework that aims to get you up and running quickly with st

382 Dec 06, 2022

PyTorch toolkit for biomedical imaging

farabio is a minimal PyTorch toolkit for out-of-the-box deep learning support in biomedical imaging. For further information, see Wikis and Docs.

47 Dec 28, 2022

Training PyTorch models with differential privacy

Opacus is a library that enables training PyTorch models with differential privacy. It supports training with minimal code changes required on the cli

1.3k Dec 29, 2022

Tutorial for surrogate gradient learning in spiking neural networks

SpyTorch A tutorial on surrogate gradient learning in spiking neural networks Version: 0.4 This repository contains tutorial files to get you started

203 Nov 28, 2022

A PyTorch implementation of Learning to learn by gradient descent by gradient descent

Intro PyTorch implementation of Learning to learn by gradient descent by gradient descent. Run python main.py TODO Initial implementation Toy data LST

300 Dec 11, 2022

You like pytorch? You like micrograd? You love tinygrad! ❤️

For something in between a pytorch and a karpathy/micrograd This may not be the best deep learning framework, but it is a deep learning framework. Due

9.7k Jan 05, 2023

PyTorch implementations of normalizing flow and its variants.

55 Dec 01, 2022

Tez is a super-simple and lightweight Trainer for PyTorch. It also comes with many utils that you can use to tackle over 90% of deep learning projects in PyTorch.

Tez: a simple pytorch trainer NOTE: Currently, we are not accepting any pull requests! All PRs will be closed. If you want a feature or something does

1.1k Jan 04, 2023

Use Jax functions in Pytorch with DLPack

Related tags

Overview

jax2torch

Install

Usage

You might also like...

A PyTorch implementation of EfficientNet

A collection of extensions and data-loaders for few-shot learning & meta-learning in PyTorch

PyTorch Extension Library of Optimized Scatter Operations

PyTorch Extension Library of Optimized Autograd Sparse Matrix Operations

Reformer, the efficient Transformer, in Pytorch

higher is a pytorch library allowing users to obtain higher order gradients over losses spanning training loops rather than individual training steps.

PyTorch implementation of TabNet paper : https://arxiv.org/pdf/1908.07442.pdf

PyTorch extensions for fast R&D prototyping and Kaggle farming

An implementation of Performer, a linear attention-based transformer, in Pytorch

Releases(0.0.7)

0.0.7(Oct 26, 2021)

0.0.6(Oct 26, 2021)

0.0.5(Oct 26, 2021)

0.0.4(Oct 26, 2021)

0.0.3(Oct 26, 2021)

0.0.2(Oct 26, 2021)

0.0.1a(Oct 26, 2021)

Owner

Phil Wang

High-level batteries-included neural network training library for Pytorch

PyTorch toolkit for biomedical imaging

Training PyTorch models with differential privacy

Tutorial for surrogate gradient learning in spiking neural networks

A PyTorch implementation of Learning to learn by gradient descent by gradient descent

You like pytorch? You like micrograd? You love tinygrad! ❤️

PyTorch implementations of normalizing flow and its variants.

Tez is a super-simple and lightweight Trainer for PyTorch. It also comes with many utils that you can use to tackle over 90% of deep learning projects in PyTorch.

Code snippets created for the PyTorch discussion board

PyGCL: Graph Contrastive Learning Library for PyTorch

A PyTorch implementation of EfficientNet

A tiny scalar-valued autograd engine and a neural net library on top of it with PyTorch-like API

A PyTorch implementation of L-BFGS.

PyTorch Extension Library of Optimized Scatter Operations

OptNet: Differentiable Optimization as a Layer in Neural Networks

Unofficial PyTorch implementation of DeepMind's Perceiver IO with PyTorch Lightning scripts for distributed training

A pure Python implementation of Compact Bilinear Pooling and Count Sketch for PyTorch.

Pretrained ConvNets for pytorch: NASNet, ResNeXt, ResNet, InceptionV4, InceptionResnetV2, Xception, DPN, etc.

GPU-accelerated PyTorch implementation of Zero-shot User Intent Detection via Capsule Neural Networks

Use Jax functions in Pytorch with DLPack