Skip to main content

Graphs, Scientific & Structured

Grids (images) and sequences (text) aren't the only structures worth modeling. This track ventures into relational data (graphs) and scientific ML, where JAX's real superpower — differentiating arbitrary functions — takes center stage.

What you'll build​

  • Graph Neural Networks (GCN) — message passing over a graph for semi-supervised node classification, using nnx.Einsum for adjacency-weighted aggregation.
  • Physics-Informed Neural Networks (PINN) — solve a differential equation by baking it into the loss, differentiating the network's output with respect to its input via jax.grad. No dataset needed.
  • Neural ODEs — continuous-depth models; learn the dynamics and integrate them through a differentiable ODE solver.
  • Tabular Deep Learning — MLPs with categorical embeddings for structured/tabular data (classification and regression).
  • Mixture of Experts (MoE) — sparse conditional computation with top-k routing and a load-balancing loss.

Prerequisites​

You should be comfortable with the training loop and custom training loops. PINNs assume basic calculus (derivatives, differential equations).

Next steps​