Discover the power of JAX in deep learning. Gain insights into its ecosystem and learn about linear algebra, pseudo-random number generation, and optimization algorithms for cleaner, structured coding.

Doxx.tar.gz

NormalCode

JuPyter

JAX is a Python library designed for high-performance ML research. It is a powerful numerical computing library, just like Numpy, but with some key improvements.

In this course, you will learn all about JAX and its ecosystem of libraries (Haiku, Jraph, Chex, Flax, Optax). Addressing a wide range of audiences, you will cover several topics including linear algebra, random variables theory, pseudo-random number generation, and optimization algorithms.

By the end of this course, you will have a new set of skills that will make deep learning programming more intuitive, structured, and clean.

Introduction to JAX and Deep Learning


## Important points

Before solving the exercises, please make sure to:

* Do _not_ change the name of the function below because it will be used by the evaluator.
* Do _not_ print anything. 
* Do _not_ import any libraries. They are pre-imported. We'll continue the convention of **`jnp`** for JAX NumPy.
* Variables are usually initialized with $1$ to help run the test case. Please **override them with intended expressions/values** in your implementation.

>**Note:** These points will also apply to all the subsequent challenges.

## Exercise 1: Testing array for zeros

Let's test our familiarity with the basic JAX and NumPy syntax.

In this simple exercise, we are given a JAX array `a`. Identify whether it contains $0$ or not. In other words, if it contains an(y) instance of zero, it should return **`true`** and vice-versa.



#a will be any JAX array
def TestZeroExists(a):
  #Please use the condition below to test for zeros
  if (1==1):  #Condition is initialized with just a dummy 1=1 case. Please override it
    return False
  else:
    return True

import jax.numpy as jnp
import jax
jax.config.update('jax_platform_name', 'cpu')




#a will be any JAX array
def TestZeroExists(a):
  if (jnp.all(a)!=0):
    return False
  else:
    return True

---

Congratulations on completing this challenge. Now we will begin the real deal by practicing Autograd.

>**Note:** For ease of coding, **`grad`** is already imported.

## Exercise 2: Autograd

Make a function $f(x,y) = x^2 - sin(y)$ and calculate its derivative with respect to both $x$ and $y$.


# Important points

Before solving the exercises, please make sure to:

* Do _not_ change the name of the function below because it will be used by the evaluator.
* Do _not_ print anything. 
* Do _not_ import any libraries. They are pre-imported. We'll continue the convention of **`jnp`** for JAX NumPy.
* Variables are usually initialized with $1$ to help run the test case. Please **override them with intended expressions/values** in your implementation.

>**Note:** These points will also apply to all the subsequent challenges.

# Exercise 1: Testing array for zeros

Let's test our familiarity with the basic JAX and NumPy syntax.

In this simple exercise, we are given a JAX array `a`. Identify whether it contains $0$ or not. In other words, if it contains an(y) instance of zero, it should return **`true`** and vice-versa.



---

Congratulations on completing this challenge. Now we will begin the real deal by practicing Autograd.

>**Note:** For ease of coding, **`grad`** is already imported.

# Exercise 2: Autograd

Make a function $f(x,y) = x^2 - sin(y)$ and calculate its derivative with respect to both $x$ and $y$.

Test your learning so far by completing a programming exercise.

Introduction

JAX Programming Model

Linear Algebra

Random Variables and Distributions

JAX Ecosystem

Project: GAN Using the JAX ecosystem

Appendix

Challenge: Basics of JAX

Important points

Exercise 1: Testing array for zeros