Skip to content

Jax Tutorial Ideas

Tutorials

  • vmap
  • jit
  • grad, jacobian, hessian
  • optimizers
  • containers - Dict, Tuple, NamedTuple

Use Cases

  • Gaussian Processes
  • Kernel Matrices and Derivatives
  • Kernel Density Estimation - Blog |
  • Optimized Kernel Ridge Regression (OKRR) - Notebook
  • Optimized Kernel Entropy Components Analysis (OKECA)
  • Centered Kernel Alignment (CKA)
  • Gaussianization Flows

Projects

  1. Uncertain Inputs for GPs

  2. Gaussianization Flows - Case Study

  3. Uncertain Inputs for Gaussian Processes

  4. Linearized GP (Taylor Expansion, 1st Order, 2nd Order)
  5. Moment Matching (RBF)
    • Scratch - Loops, numba
    • Jax
  6. MCMC Posterior Approximation
    • Scratch
    • Numpyro - NUTS/HMC
  7. Variational GP
    • Numpyro - SVI

Data

  • 1D Examples
  • 2D Examples
  • IASI Example
  • Ocean Data

Future Stuff

  • Sparse Models
  • Deep Models
  • DKL Models

Resources

GPS

OKECA