JAX topics
12 short topics in Mathematics Statistics and Computing. Each one is a single focused exercise with a clear piece of evidence at the end.
12 topics, one exercise each.
Work through them in order, or jump to the one you need.
These are draft topic outlines (draft; instructor review pending). Full lessons follow after review; ask Ajay if you need a topic sooner.
- 01
Array and tensor representation
Investigate array and tensor representation using a repeatable JAX experiment; compare results with a baseline and record limitations. Keep the exercise bounded and record the behavior supported by the chosen version.
Open micro-course → - 02
Shape and axis semantics
Investigate shape and axis semantics using a repeatable JAX experiment; compare results with a baseline and record limitations. Keep the exercise bounded and record the behavior supported by the chosen version.
Open micro-course → - 03
Broadcasting and alignment
Investigate broadcasting and alignment using a repeatable JAX experiment; compare results with a baseline and record limitations. Keep the exercise bounded and record the behavior supported by the chosen version.
Open micro-course → - 04
Numeric type and precision
Investigate numeric type and precision using a repeatable JAX experiment; compare results with a baseline and record limitations. Keep the exercise bounded and record the behavior supported by the chosen version.
Open micro-course → - 05
Vectorized operations
Investigate vectorized operations using a repeatable JAX experiment; compare results with a baseline and record limitations. Keep the exercise bounded and record the behavior supported by the chosen version.
Open micro-course → - 06
Linear algebra primitives
Investigate linear algebra primitives using a repeatable JAX experiment; compare results with a baseline and record limitations. Keep the exercise bounded and record the behavior supported by the chosen version.
Open micro-course → - 07
Statistical computations
Investigate statistical computations using a repeatable JAX experiment; compare results with a baseline and record limitations. Keep the exercise bounded and record the behavior supported by the chosen version.
Open micro-course → - 08
Random sampling and reproducibility
Investigate random sampling and reproducibility using a repeatable JAX experiment; compare results with a baseline and record limitations. Keep the exercise bounded and record the behavior supported by the chosen version.
Open micro-course → - 09
Optimization or differentiation integration
Investigate optimization or differentiation integration using a repeatable JAX experiment; compare results with a baseline and record limitations. Keep the exercise bounded and record the behavior supported by the chosen version.
Open micro-course → - 10
Memory layout and copies
Investigate memory layout and copies using a repeatable JAX experiment; compare results with a baseline and record limitations. Keep the exercise bounded and record the behavior supported by the chosen version.
Open micro-course → - 11
Numeric result validation
Investigate numeric result validation using a repeatable JAX experiment; compare results with a baseline and record limitations. Keep the exercise bounded and record the behavior supported by the chosen version.
Open micro-course → - 12
Performance and resource benchmarking
Investigate performance and resource benchmarking using a repeatable JAX experiment; compare results with a baseline and record limitations. Keep the exercise bounded and record the behavior supported by the chosen version.
Open micro-course →
Learn the same ideas beyond JAX.
These courses share topics with this one, so you can compare how each language or tool handles the same concept.
More in Data science and machine learning libraries.
You have the curiosity.
Let’s turn it into capability.
Start with a practical workshop or bring a specific question to Ajay.