JAX Pathways

HomeTake the course with you

Learner resources

Practice on your CPU, keep your own evidence, and learn in the browser, a notebook or your agent.

Run something real

Start small. Verify the result.

Each download includes the work to try and the checks to explain. Reference execution and your own exercise evidence stay separate.

Placement practice · CPU

Four logical devices

Inspect local shards, compare replication with partitioning, and repair an uneven batch on one physical host.

Open virtual-device lesson →Download notebook

Restart the notebook kernel, then Run all. Logical CPU devices do not emulate TPU performance.

Away from the browser

Read offline.

The book includes available lesson drafts, their worked code, experiments, practice and references. Unwritten roadmap topics are excluded.

Printable course readerHTML · open and print to PDFCourse EPUBFor your ebook reader

The filename retains “foundations”; both editions contain all available lesson drafts from the current build.

Your agent host

A tutor that uses this course.

Ask for a starting point, work through an exercise and keep a study plan in LEARNING.md. Skills need local course files; they do not install the course or its Python environment. Download the course workspace for the current lessons and tutors together.

  1. Choose the tools you need

    start-learning
    Choose a route or resume your plan.
    learn-jax
    Work through an authored lesson.
    jax-course-guide
    Find a topic, prerequisite or debugging lesson.
    check-jax
    Check understanding and a practical modification.

  2. Add them from your course folder

    Read the skill files first. With Node.js/npm available, run this in the extracted course folder or repository root and choose your agent when prompted:

    Terminal · course checkout
    npx skills add . --skill start-learning learn-jax jax-course-guide check-jax

    Project installation is the default. To select Codex explicitly, add --agent codex. No global installation is required.

    Install from GitHub or without Node.js

    The current tutors are not yet published on GitHub. After a revision containing these skills is published:

    Terminal · course checkout
    npx skills add mlnomadpy/jaxpathways --skill start-learning learn-jax jax-course-guide check-jax

    A private repository requires configured Git/GitHub access. Use the local command above for the current working source. The remote command becomes available once these skills are published.

    Without Node.js, download the skill bundle and copy the desired skill folders into the directory your agent host supports. The bundle also includes the contributor skill create-jax-course.

    Skills CLI installation documentation ↗ · Repository skill instructions ↗

  3. Start a conversation, then resume it

    Keep the extracted course folder or checkout open. In Codex, choose start-learning or type $start-learning. In Claude Code, use /start-learning. In another host, say “Use start-learning to help me learn JAX from this workspace.”

    The tutor keeps your goal, next lesson, attempts and review queue in LEARNING.md. Use learn-jax next time to continue. Browser-local progress is separate; neither store synchronizes automatically.

From your course checkout

Use the terminal.

Browse availability, read a lesson or run its reference companion. Execution needs the tested packages in requirements-cpu.txt.

Write your own exercise solution in the lesson script or notebook first, then run the verifier to check your assertions.

Terminal · command examples
python3 scripts/course.py list --available
python3 scripts/course.py list --pathway ship
python3 scripts/course.py show first-gradient
python3 scripts/course.py run first-gradient
python3 scripts/course.py quiz first-gradient
python3 scripts/course.py plan inference-engineer

The final command prints a proposed route. Keep the selected next lesson and your own evidence in your learning plan.

For course authors and integrations

One source, several ways to read it.

Author a lesson

Edit curriculum/course.json for membership and ordering, learning-paths/*.json for routes, and the canonical lesson.json for teaching material. Build the generated documents, notebooks, website resources and books from that source.

npm run build
npm run check
npm run smoke:cpu

Use create-jax-course to guide authoring, run npm run smoke:cpu to record CPU execution receipts, and stage accelerator benchmarks on a Cloud TPU VM.

Read-only course APIs

These are static JSON resources, with authored/planned availability. They provide no remote execution or progress service.