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

A few Windows specific scripts for PyTorch

It is a repo that contains scripts that makes using PyTorch on Windows easier. Easy Installation Update: Starting from 0.4.0, you can go to the offici

408 Dec 15, 2022

A code copied from google-research which named motion-imitation was rewrited with PyTorch

motor-system Introduction A code copied from google-research which named motion-imitation was rewrited with PyTorch. More details can get from this pr

6 Jan 08, 2022

ocaml-torch provides some ocaml bindings for the PyTorch tensor library.

ocaml-torch provides some ocaml bindings for the PyTorch tensor library. This brings to OCaml NumPy-like tensor computations with GPU acceleration and tape-based automatic differentiation.

369 Jan 03, 2023

High-fidelity performance metrics for generative models in PyTorch

5 Oct 24, 2021

A simplified framework and utilities for PyTorch

Here is Poutyne. Poutyne is a simplified framework for PyTorch and handles much of the boilerplating code needed to train neural networks. Use Poutyne

534 Dec 17, 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

Implementation of LambdaNetworks, a new approach to image recognition that reaches SOTA with less compute

Lambda Networks - Pytorch Implementation of λ Networks, a new approach to image recognition that reaches SOTA on ImageNet. The new method utilizes λ l

1.5k Jan 07, 2023

Differentiable ODE solvers with full GPU support and O(1)-memory backpropagation.

PyTorch Implementation of Differentiable ODE Solvers This library provides ordinary differential equation (ODE) solvers implemented in PyTorch. Backpr

4.4k Jan 04, 2023

pip install antialiased-cnns to improve stability and accuracy

Antialiased CNNs [Project Page] [Paper] [Talk] Making Convolutional Networks Shift-Invariant Again Richard Zhang. In ICML, 2019. Quick & easy start Ru

1.6k Dec 28, 2022

PyTorch wrappers for using your model in audacity!

130 Dec 14, 2022

Fast, general, and tested differentiable structured prediction in PyTorch

Torch-Struct: Structured Prediction Library A library of tested, GPU implementations of core structured prediction algorithms for deep learning applic

1.1k Jan 07, 2023

Pytorch implementation of Distributed Proximal Policy Optimization

Pytorch-DPPO Pytorch implementation of Distributed Proximal Policy Optimization: https://arxiv.org/abs/1707.02286 Using PPO with clip loss (from https

164 Jan 05, 2023

OptNet: Differentiable Optimization as a Layer in Neural Networks

OptNet: Differentiable Optimization as a Layer in Neural Networks This repository is by Brandon Amos and J. Zico Kolter and contains the PyTorch sourc

428 Dec 24, 2022

Code for paper "Energy-Constrained Compression for Deep Neural Networks via Weighted Sparse Projection and Layer Input Masking"

model_based_energy_constrained_compression Code for paper "Energy-Constrained Compression for Deep Neural Networks via Weighted Sparse Projection and

16 Jun 15, 2022

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

A few Windows specific scripts for PyTorch

A code copied from google-research which named motion-imitation was rewrited with PyTorch

ocaml-torch provides some ocaml bindings for the PyTorch tensor library.

High-fidelity performance metrics for generative models in PyTorch

A simplified framework and utilities for PyTorch

Tutorial for surrogate gradient learning in spiking neural networks

Implementation of LambdaNetworks, a new approach to image recognition that reaches SOTA with less compute

Differentiable ODE solvers with full GPU support and O(1)-memory backpropagation.

pip install antialiased-cnns to improve stability and accuracy

PyTorch wrappers for using your model in audacity!

Fast, general, and tested differentiable structured prediction in PyTorch

Pytorch implementation of Distributed Proximal Policy Optimization

OptNet: Differentiable Optimization as a Layer in Neural Networks

Code for paper "Energy-Constrained Compression for Deep Neural Networks via Weighted Sparse Projection and Layer Input Masking"

Distiller is an open-source Python package for neural network compression research.

PyTorch Extension Library of Optimized Scatter Operations

The goal of this library is to generate more helpful exception messages for numpy/pytorch matrix algebra expressions.

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

ONNX Runtime for PyTorch accelerates PyTorch model training using ONNX Runtime.

TorchShard is a lightweight engine for slicing a PyTorch tensor into parallel shards