Search⌘ K
Home
Courses
Introduction to JAX and Deep Learning

Vectors

Explore fundamental vector operations with JAX to strengthen your understanding of linear algebra in deep learning. Learn how to compute inner products, norms, cosine similarity, and apply Taylor approximations using JAX syntax and functions. This lesson helps you refresh key concepts while practicing efficient numerical computations essential for machine learning.

We'll cover the following...

Since we are assuming you are fairly familiar with NumPy, this chapter is by no means a thorough coverage of the topic and will just serve as a quick refresher. We will practice using JAX while discovering the equivalence between normal NumPy and the JAX version of NumPy’s syntax.

Note: As a reference, we’ll use np and jnp respectively for default and JAX NumPy versions in our codes.

Inner product

The inner product of two vectors can be calculated by any of the three syntaxes (dot(), inner() and @), as shown below:

Python 3.8
a = jnp.arange(1,10)
b = jnp.arange(11,20)
print(jnp.dot(a,b))
print(jnp.inner(a,b))
print(a@b)
print(b@a)

Linear functions

A linear function can be represented as:

f(x)=αx1+βx2+γx3+.... f(x) = \alpha x_1+\beta x_2+\gamma x_3+....

We can compose it as a lambda function too:

Python 3.8
a= 1.2
b = 2.4
FuncA = lambda x : a*x[0]+b*x[1]
x = jnp.arange(1,3)
print(FuncA(x))

Taylor approximation

Taylor approximation lies at the core of major optimization algorithms in machine learning. We can first-order approximate the value of a function near a point z as:

f^(x)=f(z)+∇f(z)T( ...