Ctrl+K

Unify code#

Unify a simple torch function and use it alongside any ML framework!

⚠️ 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.

[1]:

!pip install ivy

Firstly, let’s import the dependencies and define a torch function.

[ ]:

import ivy
import torch

def normalize(x):
    mean = torch.mean(x)
    std = torch.std(x)
    return torch.div(torch.sub(x, mean), std)

By using ivy.unify(), you can convert any code from any framework into Ivy code, which as we have already seen, can be executed using any framework as the backend.

Let’s unify the function!

[3]:

normalize = ivy.unify(normalize, source="torch", with_numpy=False)

And that’s it! The normalize function can now be used with any ML framework. It’s as simple as that!

So, let’s give it a try!

[4]:

# import the frameworks
import numpy as np
import jax.numpy as jnp
import tensorflow as tf

[5]:

# create random numpy arrays for testing
np.random.seed(123) # for reproducibility
x = np.random.uniform(size=10).astype(np.float32)

ivy.set_backend("numpy")
print(normalize(x))

# jax
x_ = jnp.array(x)
ivy.set_backend("jax")
print(normalize(x_))

# tensorflow
x_ = tf.constant(x)
ivy.set_backend("tensorflow")
print(normalize(x_))

# torch
x_ = torch.tensor(x)
ivy.set_backend("torch")
print(normalize(x_))

ivy.array([ 0.66391283, -1.12516928, -1.38367081,  0.03102401,  0.76419425,
       -0.52797794,  1.90346956,  0.61316347, -0.27585283, -0.66309303])

/workspaces/tracer-transpiler/ivy_repo/ivy/utils/exceptions.py:383: 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(

ivy.array([ 0.66391283, -1.12516928, -1.38367081,  0.03102401,  0.76419425,
       -0.52797794,  1.90346956,  0.61316347, -0.27585283, -0.66309303])

/workspaces/tracer-transpiler/ivy_repo/ivy/utils/exceptions.py:383: UserWarning: The current backend: 'tensorflow' 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(

ivy.array([ 0.66391283, -1.12516928, -1.38367081,  0.03102401,  0.76419425,
       -0.52797794,  1.90346956,  0.61316347, -0.27585283, -0.66309303])
ivy.array([ 0.66391283, -1.12516928, -1.38367081,  0.03102401,  0.76419425,
       -0.52797794,  1.90346956,  0.61316347, -0.27585283, -0.66309303])

We can see that the new normalize function can operate with any ML framework. ivy.unify converts the framework-specific torch implementation into a framework-agnostic ivy implementation, which is compatible with all frameworks.

Round Up#

That’s it, you can now unify ML code! However, there are several other important topics to master before you’re ready to unify ML code like a pro 🥷. Next, we’ll be learning how to make our unified Ivy code run much more efficiently! ⚡

On this page