Add intro to JAX

[takumiohym]

Added the solution notebook for introduction to JAX.

I'll add lab version after initial review.

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[gemini-code-assist]

Summary of Changes

Hello @takumiohym, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request introduces a new Jupyter notebook aimed at providing a foundational understanding of JAX. The notebook covers essential JAX functionalities, such as array operations and its functional programming paradigm, and demonstrates its automatic differentiation capabilities through a practical linear regression example. It serves as a hands-on guide for users to grasp JAX's core features and their application in machine learning.

Highlights

• New JAX Introduction Notebook: A new Jupyter notebook has been added, providing a comprehensive introduction to the JAX library.
• Core JAX Concepts Covered: The notebook familiarizes users with fundamental JAX operations, including array manipulation, understanding immutability, and seamless interoperability with NumPy.
• Automatic Differentiation Demonstrated: It showcases JAX's powerful jax.grad for automatic differentiation, a crucial feature for machine learning model training.
• Linear Regression Implementation: Users are guided through implementing a linear regression model from scratch, covering loss function definition, gradient computation, and a basic training loop using JAX.
• Feature Engineering Example: A bonus section illustrates how to tackle non-linear functions by applying feature engineering within the JAX framework.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request adds a solution notebook for an introduction to JAX. The notebook is well-structured and covers the basics of JAX, including array operations, automatic differentiation, and training a linear regression model. My review includes a critical fix for data leakage in the baseline model calculation, along with several suggestions to improve code clarity and reduce redundancy, such as removing duplicate variable definitions and renaming functions that are redefined. Also, please consider moving the notebook to a more appropriate directory, as it's currently under introduction_to_tensorflow.

[takumiohym]

Rebased onto new_env_test.