Python tutorials

A learning path through eagle from Python. Start with the Python quickstart, then follow the groups in order. Each tutorial is a short notebook that runs on the CPU and picks up the GPU automatically when one is visible. The C++ path is on the C++ tutorials page.

Launching and graphs

What a graph is, how many steps to fuse into one launch, how finished samples leave the run, and where the work runs.

• What a graph is: one launch replays a whole run.

• Steps per launch: fuse steps to cut host-device chatter.

• Finished samples: compaction, reorder, and when not to.

• A multi-kernel step: propagate, diagnose and stop, declaratively.

• Host or device?: let the data decide where it runs.

Training and gradients

Putting a kernel inside a training loop, and differentiating a whole graph.

• Gradients through a physics kernel: a hawk kernel inside loss.backward().

• Build a step as a graph: eagle.Graph, built, inspected and run.

• Gradients through a graph: eagle.diff, checked against finite differences.

• Conditionals, budgets and events: reverse through a skipped step, a capped memory budget, and a stop between two steps.

• A hawk kernel inside your jitted JAX step: fit a drag coefficient with jax.jit and grad.

• Batches and Jacobians: vmap, jacfwd, hessian, and which cost one launch.

• A stepping loop in lax.scan: balls that land and stop, on the CPU and the GPU.

• Train a torch net inside your simulation loop: a net runs before every step; the gradient goes through the whole run.

• A Warp kernel in the middle of an eagle graph: your kernel between two eagle steps, with its gradient.

• A JAX model between eagle launches: your network writes the correction, eagle steps the batch.

• The eagle batch inside jit: while_loop to run, scan to differentiate, no where.

Seeing what ran

• Profile your run: where the time went, in three numbers.

Further

• A graph with parallel branches: fork, join and reduce, optional reading.