Skip to content

Commit

Permalink
Add readme
Browse files Browse the repository at this point in the history
  • Loading branch information
takuseno committed Oct 25, 2024
1 parent 7179c2d commit 8045f26
Showing 1 changed file with 43 additions and 0 deletions.
43 changes: 43 additions & 0 deletions README.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,43 @@
[![test](https://github.com/takuseno/tensor-bridge/actions/workflows/test.yaml/badge.svg)](https://github.com/takuseno/tensor-bridge/actions/workflows/test.yaml)

# tensor-bridge
`tensor-bridge` is a light-weight library that achieves inter-library tensor transfer by native `cudaMemcpy` call for minimal overheads.

```py
import torch
import jax
from tensor_bridge import copy_tensor


# PyTorch tensor
torch_data = torch.rand(2, 3, 4, device="cuda:0")

# Jax tensor
jax_data = jax.random.uniform(jax.random.key(123), shape=(2, 3, 4))

# Copy PyTorch tensor to Jax tensor
copy_tensor(torch_data, jax_data)

# And, other way around
copy_tensor(jax_data, torch_data)
```

## Supported deep learning libraries
- PyTorch
- Jax


## Installation
Your macine needs to install `nvcc` to compile a native code.
```
pip install git+https://github.com/takuseno/tensor-bridge
```
Pre-built package release is in progress.


## Unit test
You machine needs to install NVIDIA's GPU and nvidia-driver to execute tests.
```
./bin/build-docker
./bin/test
```

0 comments on commit 8045f26

Please sign in to comment.