ETH Zürich AISE: Introduction to JAX

↓↓↓ LECTURE OVERVIEW BELOW ↓↓↓
ETH Zürich AI in the Sciences and Engineering 2024

Course Website (links to slides and tutorials): https://www.camlab.ethz.ch/teaching/a...

Lecturers: Dr. Ben Moseley and Prof. Siddhartha Mishra

▬ Lecture Content ▬▬▬
0:00 - Introduction
2:02 - What is JAX?
3:08 - JAX in ML and scientific computing
5:34 - Accelerated array computation
14:45 - Example: wave simulation with JAX
19:09 - Program transformation
28:28 - 🔴 Live coding: autodiff in JAX | Code: https://github.com/benmoseley/AISE-2024
32:45 - Advanced autodiff
33:55 - Automatic vectorisation
42:38 - 🔴 Vectorising a layer function
44:59 - Just-in-time (JIT) compilation
49:32 - 🔴 Measuring JIT speed-up
50:20 - 🔴 Putting it all together: linear regression
57:00 - JAX ecosystem
1:01:28 - Example: optimisation with JAX
1:04:34 - Summary

▬ Course Overview ▬▬▬
Lecture 1: Course Introduction    • ETH Zürich AISE: Course Introduction  
Lecture 2: Introduction to Deep Learning Part 1    • ETH Zürich AISE: Introduction to Deep Lear...  
Lecture 3: Introduction to Deep Learning Part 2    • ETH Zürich AISE: Introduction to Deep Lear...  
Lecture 4: Importance of PDEs in Science    • ETH Zürich AISE: Importance of PDEs in Sci...  
Lecture 5: Physics-Informed Neural Networks – Introduction    • ETH Zürich AISE: Physics-Informed Neural N...  
Lecture 6: Physics-Informed Neural Networks – Limitations and Extensions Part 1    • ETH Zürich AISE: Physics-Informed Neural N...  
Lecture 7: Physics-Informed Neural Networks – Limitations and Extensions Part 2    • ETH Zürich AISE: Physics-Informed Neural N...  
Lecture 8: Physics-Informed Neural Networks – Theory Part 1    • ETH Zürich AISE: Physics-Informed Neural N...  
Lecture 9: Physics-Informed Neural Networks – Theory Part 2    • ETH Zürich AISE: Physics-Informed Neural N...  
Lecture 10: Introduction to Operator Learning Part 1    • ETH Zürich AISE: Introduction to Operator ...  
Lecture 11: Introduction to Operator Learning Part 2    • ETH Zürich AISE: Introduction to Operator ...  
Lecture 12: Fourier Neural Operators    • ETH Zürich AISE: Fourier Neural Operators  
Lecture 13: Spectral Neural Operators and Deep Operator Networks    • ETH Zürich AISE: Spectral Neural Operators...  
Lecture 14: Convolutional Neural Operators    • ETH Zürich AISE: Convolutional Neural Oper...  
Lecture 15: Time-Dependent Neural Operators    • ETH Zürich AISE: Time-Dependent Neural Ope...  
Lecture 16: Large-Scale Neural Operators    • ETH Zürich AISE: Large-Scale Neural Operators  
Lecture 17: Attention as a Neural Operator    • ETH Zürich AISE: Attention as a Neural Ope...  
Lecture 18: Windowed Attention and Scaling Laws    • ETH Zürich AISE: Windowed Attention and Sc...  
Lecture 19: Introduction to Hybrid Workflows Part 1    • ETH Zürich AISE: Introduction to Hybrid Wo...  
Lecture 20: Introduction to Hybrid Workflows Part 2    • ETH Zürich AISE: Introduction to Hybrid Wo...  
Lecture 21: Neural Differential Equations    • ETH Zürich AISE: Neural Differential Equat...  
Lecture 22: Introduction to Diffusion Models    • ETH Zürich AISE: Introduction to Diffusion...  
Lecture 23: Introduction to JAX    • ETH Zürich AISE: Introduction to JAX  
Lecture 24: Symbolic Regression and Model Discovery    • ETH Zürich AISE: Symbolic Regression and M...  
Lecture 25: Applications of AI in Chemistry and Biology Part 1    • ETH Zürich AISE: Applications of AI in Che...  
Lecture 26: Applications of AI in Chemistry and Biology Part 2    • ETH Zürich AISE: Applications of AI in Che...  

▬ Course Description ▬▬▬
AI is having a profound impact on science by accelerating discoveries across physics, chemistry, biology, and engineering. This course presents a highly topical selection of AI applications across these fields. Emphasis is placed on using AI, particularly deep learning, to understand systems modelled by PDEs, and key scientific machine learning concepts and themes are discussed.

▬ Course Learning Objectives ▬▬▬
Aware of advanced applications of AI in the sciences and engineering
Familiar with the design, implementation, and theory of these algorithms
Understand the pros/cons of using AI and deep learning for science
Understand key scientific machine learning concepts and themes