"""
<center>
    <h1>JaxTon</h1>
    <i>💯 JAX exercises</i>
    <br>
    <br>
    <a href='https://github.com/vopani/jaxton/blob/master/LICENSE'>
        <img src='https://img.shields.io/badge/license-Apache%202.0-blue.svg?logo=apache'>
    </a>
    <a href='https://github.com/vopani/jaxton'>
        <img src='https://img.shields.io/github/stars/vopani/jaxton?color=yellowgreen&logo=github'>
    </a>
    <a href='https://twitter.com/vopani'>
        <img src='https://img.shields.io/twitter/follow/vopani'>
    </a>
</center>
"""
"""
<center>
    This is Set 1: JAX Introduction (Exercises 1-10) of <b>JaxTon</b>: <i>💯 JAX exercises</i>
    <br>
    You can find all the exercises and solutions on <a href="https://github.com/vopani/jaxton#exercises-">GitHub</a>
</center>
"""
"""
**Prerequisites**

* The configuration of jax should be updated as shown in the code snippet below in order to use TPUs
"""
## install jax
!python3 -m pip install jax
## import packages
import jax
import os
import requests

## setup JAX to use TPUs if available
try:
    url = 'http:' + os.environ['TPU_NAME'].split(':')[1] + ':8475/requestversion/tpu_driver_nightly'
    resp = requests.post(url)
    jax.config.FLAGS.jax_xla_backend = 'tpu_driver'
    jax.config.FLAGS.jax_backend_target = os.environ['TPU_NAME']
except:
    pass

jax.devices()
"""
**Exercise 1: Install the `jax` package**
"""
"""
**Exercise 2: Import the `jax` package**
"""
"""
**Exercise 3: Display the version of `jax`**
"""
"""
**Exercise 4: Display the default backend of `jax`**
"""
"""
**Exercise 5: Display the devices of the backend**
"""
"""
**Exercise 6: Create a JAX DeviceArray with values [10, 1, 24] and assign it to `data`**
"""
"""
**Exercise 7: Display the type of `data`**
"""
"""
**Exercise 8: Display the shape of `data`**
"""
"""
**Exercise 9: Transfer `data` to host and assign it to `data_host`**
"""
"""
**Exercise 10: Transfer `data_host` to device and assign it to `data_device`**
"""
"""
<center>
    This completes Set 1: JAX Introduction (Exercises 1-10) of <b>JaxTon</b>: <i>💯 JAX exercises</i>
    <br>
    You can find all the exercises and solutions on <a href="https://github.com/vopani/jaxton#exercises-">GitHub</a>
</center>
"""