Get Started with JAX A.I. Stack

Join this workshop if you're interested in getting started with JAX for A.I. research or scientific computing. JAX is an open-source machine learning library used by Google to create and train our foundational models, and for most of our A.I. research across the company. It's popular with frontier model developers, and machine learning researchers. In this workshop, we'll help you get started with JAX, using a mix of slides and hands-on coding exercises. We'll assume you're new to both JAX and its ecosystem of libraries (NNX, Orbax, Grain, etc). Our goal is to give you the flavor of these to help you get started, and so know where to go to learn more.

Please bring a laptop. We'll use Google Colab to run our example. There is nothing to install in advance.

The speaker will be around afterwards and will be happy to chat with you 1:1 about research opportunities and collaborations.

Knowledge prerequisites: You should be familiar with Python programming, and have taken at least one machine learning course.

Hardware/software prerequisites: Bring a laptop with a web browser.

Workshop format: Demonstration and hands-on

Target audience: Students, researchers, faculty, staff

See the full Research Computing training schedule or subscribe to the Research Computing mailing list.

Research Computing | View More Events
Co-hosted with: GradFUTURES

Contact the organizers