Jax Vs PyTorch [Key Differences]

In this Python tutorial, we will learn about Jax Vs PyTorch in python. Jax is a programming language showing and creating transformations of numerical programs. It is also able to compile numerical programs for CPU or accelerating GPU.

And PyTorch is an open-source machine learning library that is mostly used for computer vision and natural language processing in python. And we will also cover different examples related to Jax vs PyTorch. And these are the following topics that we are going to discuss in this tutorial.

• Introduction to Jax
• Introduction to PyTorch
• Jax Vs PyTorch
• Jax Vs PyTorch Vs Tensorflow
• Jax Vs PyTorch benchmark

Introduction to Jax

In this section, we will learn about What is JAX and how its works in python.

JAX stands for Just After execution. It is a machine learning library developed by DeepMind. Jax is a JIT( Just In Time ) compiler that focused on controlling the maximum numbers of FLOPS that create optimized code while using python.

Jax is a programming language showing and creating transformations of numerical programs. It is also able to compile numerical programs for CPU or accelerating GPU.

Jax enables Numpy code ON not only for CPU but GPU and TPU as well.

Code:

In the following code, we will import all the necessary libraries such as import jax.numpy as jnp, import grad, jit, vmap from jax, and import random from jax.

• rd = random.PRNGKey(0) is used to generate random data and the random state is described by two unsigned 32-bit integers that we call as a key.
• y = random.normal(rd, (8,)) is used to generate a sample of numbers drawn from the normal distribution.
• print(y) is used to print the y values using the print() function.
• size = 2899 is used to give the size.
• random.normal(rd, (size, size), dtype=jnp.float32) is used to generate a sample of numbers drawn from the normal distribution.
• %timeit jnp.dot(y, y.T).block_until_ready() it is runs on the GPU when GPU is available otherwise it is runs on the CPU.

# Importing libraries
import jax.numpy as jnp
from jax import grad, jit, vmap
from jax import random
# Generating random data
rd = random.PRNGKey(0)
y = random.normal(rd, (8,))
print(y)
# Multiply two matrices
size = 2899
y = random.normal(rd, (size, size), dtype=jnp.float32)
# runs on the GPU
%timeit jnp.dot(y, y.T).block_until_ready()  

Output:

After running the above code, we get the following output in which we can see that the multiplication of the two matrices is printed on the screen.

So with this, we understood Jax and how its works in python.

Read: PyTorch Activation Function

Introduction to PyTorch

In this section, we will learn about What is PyTorch and how we can work with PyTorch in python.

PyTorch is an open-source machine learning library that is mostly used for computer vision and natural language processing in python. It is developed by the Facebook AI(Artificial Intelligence) Research Lab.

It is a software released under the Modified BSD License. It is built based on python which supports the calculation of tensors on GPU ( Graphical Processing Unit).

PyTorch is easy to use, has efficient memory usage, dynamic computational graph, is flexible, and creates coding feasible that increases the processing speed. The PyTorch is the most recommended library for deep learning and artificial intelligence.

Code:

In the following code, we will import all the necessary libraries such as import torch and import math.

• y = torch.linspace(-math.pi, math.pi, 2000, device=device, dtype=dtype) is used to create random input and output data.
• m = torch.randn((), device=device, dtype=dtype) is used to randomly initialize the weights.
• z_pred = m + n * y + o * y ** 2 + p * y ** 3 is used as a forward pass to compute predicted Z.
• loss = (z_pred – z).pow(2).sum().item() is used to compute the loss.
• print(int, loss) is used to print the loss.
• grad_m = grad_z_pred.sum(): Here we apply backpropagation to compute gradients of m,n,o, and p concerning loss.
• m -= learning_rate * grad_m is used to update weights using gradient descent.
• print(f’Result: z = {m.item()} + {n.item()} y + {o.item()} y^2 + {p.item()} y^3′) is used to print the result using print() function.

# Importing libraries
import torch 
import math 

# Device Used
dtype = torch.float 
device = torch.device("cpu") 

# Create random input and output data 
y = torch.linspace(-math.pi, math.pi, 2000, device=device, dtype=dtype) 
z = torch.sin(y)  

# Randomly initialize weights 
m = torch.randn((), device=device, dtype=dtype) 
n = torch.randn((), device=device, dtype=dtype) 
o = torch.randn((), device=device, dtype=dtype) 
p = torch.randn((), device=device, dtype=dtype) 


learning_rate = 1e-6 
for i in range(2000): 
    # Forward pass: compute predicted z 
    z_pred = m + n * y + o * y ** 2 + p * y ** 3 
 
    # Compute and print loss 
    loss = (z_pred - z).pow(2).sum().item() 
    if i % 100 == 99: 
        print(int, loss) 

    # Backprop to compute gradients of m, n, o, p with respect to loss 
    grad_z_pred = 2.0 * (z_pred - z) 
    grad_m = grad_z_pred.sum() 
    grad_n = (grad_z_pred * y).sum() 
    grad_o = (grad_z_pred * y ** 2).sum() 
    grad_p = (grad_z_pred * y ** 3).sum() 
 
    # Update weights using gradient descent 
    m -= learning_rate * grad_m 
    n -= learning_rate * grad_n 
    o -= learning_rate * grad_o 
    p -= learning_rate * grad_p 
 
# Print the result
print(f'Result: z = {m.item()} + {n.item()} y + {o.item()} y^2 + {p.item()} y^3') 

Output:

In the below output, you can see that the result of the items is printed on the screen.

So, with this, we understood PyTorch and its implementation.

Read: PyTorch fully connected layer

Jax Vs PyTorch

In this section, we will learn about the Key differences between Jax and PyTorch in python.

• Introduction to Jax
• Introduction to PyTorch
• Jax Vs PyTorch
• Jax Vs PyTorch Vs TensorFlow
• Jax Vs PyTorch benchmark