Gain insights into JAX and Flax's features for deep learning. Learn about optimizers, functions, data loading, and model training. Explore hands-on projects for practical experience.

edujax_1.tar.gz

jaxbook

jax-code-widget

This course comprehensively introduces JAX and Flax, two open-source libraries that have gained prominence for their efficiency, flexibility, and scalability in deep learning applications. 

In this course, you’ll explore deep learning principles and understand the unique features of JAX and Flax. You will learn the basics of JAX, optimizers using JAX and Flax, and loss and activation functions. You’ll also learn how to load datasets, perform classification using distributed learning, and use ResNet and LSTM models. In the end, you will complete a project for hands-on experience using JAX and Flax for transfer learning. 

By the end of the course, you’ll be proficient in implementing and customizing neural network models using JAX and Flax, equipped with hands-on skills in advanced optimization and distributed training.

Deep Learning with JAX and Flax

**(True or False)** JAX provides functionalities similar to NumPy. 

(True or False) JAX provides functionalities similar to NumPy.

**(True or False)** Generating random numbers in JAX and NumPy is similar.

(True or False) Generating random numbers in JAX and NumPy is similar.

**(True or False)** The `print()` function introduces side effects in JAX functions. 

(True or False) The <code>print()</code> function introduces side effects in JAX functions.

**(Select all that apply.)** Which are the ways of creating arrays in JAX? 

(Select all that apply.) Which are the ways of creating arrays in JAX?

**(True or False)** JAX arrays can be modified. 

(True or False) JAX arrays can be modified.

What are the values of variables `a` and `b` in the following snippet:
```python
import jax
seed = 98
key = jax.random.PRNGKey(seed)
jax.random.uniform(key)
key, subkey = jax.random.split(key)
a = jax.random.uniform(subkey)
b = jax.random.uniform(subkey)
```

What are the values of variables <code>a</code> and <code>b</code> in the following snippet:
<pre><code class="language-python">import jax
seed = 98
key = jax.random.PRNGKey(seed)
jax.random.uniform(key)
key, subkey = jax.random.split(key)
a = jax.random.uniform(subkey)
b = jax.random.uniform(subkey)
</code></pre>

Consider the following code. How many times will the statement `"Executing function"` be printed?
```python
def impure_print_side_effect(x):
    print("Executing function")
    return x
print("First call:", jax.jt(impure_print_side_effect)(10.))
print("Second call:", jax.jt(impure_print_side_effect)(100.))
print("Third call:", jax.jt(impure_print_side_effect)(jnp.array([111.]))) 
```

Consider the following code. How many times will the statement <code>&quot;Executing function&quot;</code> be printed?
<pre><code class="language-python">def impure_print_side_effect(x):
 print(&quot;Executing function&quot;)
 return x
print(&quot;First call:&quot;, jax.jt(impure_print_side_effect)(10.))
print(&quot;Second call:&quot;, jax.jt(impure_print_side_effect)(100.))
print(&quot;Third call:&quot;, jax.jt(impure_print_side_effect)(jnp.array([111.]))) 
</code></pre>

What is the output of the following code: `jnp.sum([1, 2, 3])`?

What is the output of the following code: <code>jnp.sum([1, 2, 3])</code>?

Test the concepts you have learned so far.

Course Introduction

Basics of JAX

Optimizers in JAX and Flax

Loss and Activation Functions

Load Datasets in JAX

Image Classification and Distributed Training

TensorBoard and State Handling

LSTM in JAX and Flax

Flax vs. TensorFlow

Using ResNet Model in Flax

Transfer Learning in JAX and Flax

Conclusion

Appendix

Quiz: Basics of JAX