Multi-Layer Perceptron
Physics-Informed Neural Networks (PINNs)
DeepONet: Operator Learning with Neural Networks
Differentiable Simulation
Graph Neural Network
Lab Setup

JAX: The Scientific Computing Advantage

JAX provides several advantages for scientific computing:

  • Functional programming: Pure functions, no side effects
  • JIT compilation: Fast execution with jit
  • Vectorization: Automatic batching with vmap
  • Transformations: grad, jit, vmap compose seamlessly

Same computation in JAX:

Back
 
Next
©  |   Cornell University    |   Center for Advanced Computing    |   Copyright Statement    |   Access Statement
CVW material development is supported by NSF OAC awards 1854828, 2321040, 2323116 (UT Austin) and 2005506 (Indiana University)