Video Tutorial

Transpile any library

Transpile the kornia library to jax with just one line of code.

⚠️ If you are running this notebook in Colab, you will have to install Ivy and some dependencies manually. You can do so by running the cell below ⬇️

If you want to run the notebook locally but don’t have Ivy installed just yet, you can check out the Get Started section of the docs.

!pip install ivy
!pip install kornia

Requirement already satisfied: ivy in /workspaces/ivy (0.0.4.0)
Requirement already satisfied: numpy in /opt/fw/mxnet (from ivy) (1.26.1)
Requirement already satisfied: einops in /opt/miniconda/envs/multienv/lib/python3.10/site-packages (from ivy) (0.7.0)
Requirement already satisfied: psutil in /opt/miniconda/envs/multienv/lib/python3.10/site-packages (from ivy) (5.9.6)
Requirement already satisfied: termcolor in /opt/fw/tensorflow (from ivy) (2.3.0)
Requirement already satisfied: colorama in /opt/miniconda/envs/multienv/lib/python3.10/site-packages (from ivy) (0.4.6)
Requirement already satisfied: packaging in /opt/fw/tensorflow (from ivy) (23.2)
Requirement already satisfied: nvidia-ml-py in /opt/miniconda/envs/multienv/lib/python3.10/site-packages (from ivy) (12.535.108)
Requirement already satisfied: diskcache in /opt/miniconda/envs/multienv/lib/python3.10/site-packages (from ivy) (5.6.3)
Requirement already satisfied: google-auth in /opt/fw/tensorflow (from ivy) (2.23.3)
Requirement already satisfied: urllib3<2.0 in /root/.local/lib/python3.10/site-packages (from ivy) (1.26.18)
Requirement already satisfied: requests in /opt/fw/mxnet (from ivy) (2.31.0)
Requirement already satisfied: pyvis in /opt/miniconda/envs/multienv/lib/python3.10/site-packages (from ivy) (0.3.2)
Requirement already satisfied: dill in /opt/miniconda/envs/multienv/lib/python3.10/site-packages (from ivy) (0.3.7)
Requirement already satisfied: astunparse in /opt/fw/tensorflow (from ivy) (1.6.3)
Requirement already satisfied: ml-dtypes in /opt/fw/tensorflow (from ivy) (0.2.0)
Requirement already satisfied: cloudpickle in /opt/fw/tensorflow (from ivy) (3.0.0)
Requirement already satisfied: gast in /opt/fw/tensorflow (from ivy) (0.5.4)
Requirement already satisfied: tqdm in /opt/miniconda/envs/multienv/lib/python3.10/site-packages (from ivy) (4.66.1)
Requirement already satisfied: wheel<1.0,>=0.23.0 in /opt/fw/tensorflow (from astunparse->ivy) (0.41.3)
Requirement already satisfied: six<2.0,>=1.6.1 in /opt/fw/tensorflow (from astunparse->ivy) (1.16.0)
Requirement already satisfied: cachetools<6.0,>=2.0.0 in /opt/fw/tensorflow (from google-auth->ivy) (5.3.2)
Requirement already satisfied: pyasn1-modules>=0.2.1 in /opt/fw/tensorflow (from google-auth->ivy) (0.3.0)
Requirement already satisfied: rsa<5,>=3.1.4 in /opt/fw/tensorflow (from google-auth->ivy) (4.9)
Requirement already satisfied: ipython>=5.3.0 in /opt/miniconda/envs/multienv/lib/python3.10/site-packages (from pyvis->ivy) (8.17.1)
Requirement already satisfied: jinja2>=2.9.6 in /opt/fw/torch (from pyvis->ivy) (3.1.2)
Requirement already satisfied: jsonpickle>=1.4.1 in /opt/miniconda/envs/multienv/lib/python3.10/site-packages (from pyvis->ivy) (3.0.2)
Requirement already satisfied: networkx>=1.11 in /opt/fw/torch (from pyvis->ivy) (3.2.1)
Requirement already satisfied: charset-normalizer<4,>=2 in /opt/fw/mxnet (from requests->ivy) (3.3.1)
Requirement already satisfied: idna<4,>=2.5 in /opt/fw/mxnet (from requests->ivy) (3.4)
Requirement already satisfied: certifi>=2017.4.17 in /opt/fw/mxnet (from requests->ivy) (2023.7.22)
Requirement already satisfied: decorator in /opt/fw/tensorflow (from ipython>=5.3.0->pyvis->ivy) (5.1.1)
Requirement already satisfied: jedi>=0.16 in /opt/miniconda/envs/multienv/lib/python3.10/site-packages (from ipython>=5.3.0->pyvis->ivy) (0.19.1)
Requirement already satisfied: matplotlib-inline in /opt/miniconda/envs/multienv/lib/python3.10/site-packages (from ipython>=5.3.0->pyvis->ivy) (0.1.6)
Requirement already satisfied: prompt-toolkit!=3.0.37,<3.1.0,>=3.0.30 in /opt/miniconda/envs/multienv/lib/python3.10/site-packages (from ipython>=5.3.0->pyvis->ivy) (3.0.39)
Requirement already satisfied: pygments>=2.4.0 in /opt/fw/jax (from ipython>=5.3.0->pyvis->ivy) (2.16.1)
Requirement already satisfied: stack-data in /opt/miniconda/envs/multienv/lib/python3.10/site-packages (from ipython>=5.3.0->pyvis->ivy) (0.6.3)
Requirement already satisfied: traitlets>=5 in /opt/miniconda/envs/multienv/lib/python3.10/site-packages (from ipython>=5.3.0->pyvis->ivy) (5.13.0)
Requirement already satisfied: exceptiongroup in /opt/fw/paddle (from ipython>=5.3.0->pyvis->ivy) (1.1.3)
Requirement already satisfied: pexpect>4.3 in /opt/miniconda/envs/multienv/lib/python3.10/site-packages (from ipython>=5.3.0->pyvis->ivy) (4.8.0)
Requirement already satisfied: MarkupSafe>=2.0 in /opt/fw/tensorflow (from jinja2>=2.9.6->pyvis->ivy) (2.1.3)
Requirement already satisfied: pyasn1<0.6.0,>=0.4.6 in /opt/fw/tensorflow (from pyasn1-modules>=0.2.1->google-auth->ivy) (0.5.0)
Requirement already satisfied: parso<0.9.0,>=0.8.3 in /opt/miniconda/envs/multienv/lib/python3.10/site-packages (from jedi>=0.16->ipython>=5.3.0->pyvis->ivy) (0.8.3)
Requirement already satisfied: ptyprocess>=0.5 in /opt/miniconda/envs/multienv/lib/python3.10/site-packages (from pexpect>4.3->ipython>=5.3.0->pyvis->ivy) (0.7.0)
Requirement already satisfied: wcwidth in /opt/miniconda/envs/multienv/lib/python3.10/site-packages (from prompt-toolkit!=3.0.37,<3.1.0,>=3.0.30->ipython>=5.3.0->pyvis->ivy) (0.2.9)
Requirement already satisfied: executing>=1.2.0 in /opt/miniconda/envs/multienv/lib/python3.10/site-packages (from stack-data->ipython>=5.3.0->pyvis->ivy) (2.0.1)
Requirement already satisfied: asttokens>=2.1.0 in /opt/miniconda/envs/multienv/lib/python3.10/site-packages (from stack-data->ipython>=5.3.0->pyvis->ivy) (2.4.1)
Requirement already satisfied: pure-eval in /opt/miniconda/envs/multienv/lib/python3.10/site-packages (from stack-data->ipython>=5.3.0->pyvis->ivy) (0.2.2)
WARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv
Requirement already satisfied: kornia in /opt/miniconda/envs/multienv/lib/python3.10/site-packages (0.7.0)
Requirement already satisfied: packaging in /opt/fw/tensorflow (from kornia) (23.2)
Requirement already satisfied: torch>=1.9.1 in /opt/fw/torch (from kornia) (2.1.0+cpu)
Requirement already satisfied: filelock in /opt/fw/torch (from torch>=1.9.1->kornia) (3.13.1)
Requirement already satisfied: typing-extensions in /opt/fw/tensorflow (from torch>=1.9.1->kornia) (4.8.0)
Requirement already satisfied: sympy in /opt/fw/torch (from torch>=1.9.1->kornia) (1.12)
Requirement already satisfied: networkx in /opt/fw/torch (from torch>=1.9.1->kornia) (3.2.1)
Requirement already satisfied: jinja2 in /opt/fw/torch (from torch>=1.9.1->kornia) (3.1.2)
Requirement already satisfied: fsspec in /opt/fw/jax (from torch>=1.9.1->kornia) (2023.10.0)
Requirement already satisfied: MarkupSafe>=2.0 in /opt/fw/tensorflow (from jinja2->torch>=1.9.1->kornia) (2.1.3)
Requirement already satisfied: mpmath>=0.19 in /opt/fw/torch (from sympy->torch>=1.9.1->kornia) (1.3.0)
WARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv

In previous tutorials, we demonstrated how to transpile simple functions from one framework to another using ivy.transpile. However, in real-world scenarios, you often need to access all functions from a specific library. Fortunately, the transpiler is capable of doing just that. Let’s explore a simple example where we convert the kornia library from torch to jax.

First, let’s import everything we need:

import ivy
import kornia
import requests
import jax
import jax.numpy as jnp
import numpy as np
from PIL import Image

Now we can transpile the library to jax. Since it’s not practical to specify arguments for every function, we’ll transpile it lazily.

jax.config.update('jax_enable_x64', True)
jax_kornia = ivy.transpile(kornia, source="torch", to="jax")

WARNING:root:To preserve the tracer and transpiler caches across multiple machines, ensure that the relative path of your projects from the .ivy folder is consistent across all machines. You can do this by adding .ivy to your home folder and placing all projects in the same place relative to the home folder on all machines.
/workspaces/ivy/ivy/utils/exceptions.py:390: UserWarning: The current backend: 'jax' does not support inplace updates natively. Ivy would quietly create new arrays when using inplace updates with this backend, leading to memory overhead (same applies for views). If you want to control your memory management, consider doing ivy.set_inplace_mode('strict') which should raise an error whenever an inplace update is attempted with this backend.
  warnings.warn(

Let’s load a sample image and convert it to the format expected by kornia. Keep in mind that even though the operations will be performed in jax, the transpiler traces a computational graph, so we still need to use kornia’s data format.

url = "http://images.cocodataset.org/train2017/000000000034.jpg"
raw_img = Image.open(requests.get(url, stream=True).raw)
img = jnp.transpose(jnp.array(raw_img), (2, 0, 1))
img = jnp.expand_dims(img, 0) / 255
display(raw_img)

Now that we have our sample image, we can easily call any kornia function using our transpiled version of the library, jax_kornia. As expected, both inputs and outputs of this function are jax.Array instances.

out = jax_kornia.enhance.sharpness(img, 10)
type(out)

jaxlib.xla_extension.ArrayImpl

Finally, we can verify that the transformation has been applied correctly!

np_image = np.uint8(np.array(out[0])*255)
display(Image.fromarray(np.transpose(np_image, (1, 2, 0))))

Round Up

Congratulations! 🎉 You are now capable of using any array computing library in your preferred framework leveraging ivy.transpile. In the next tutorial, we will explore how to convert trainable modules and layers from one framework to another ➡️