AI/ML for Scientific Discovery Using PyTorch and JAX
Build intelligent models for data-driven scientific research using PyTorch and JAX.
Skills you will gain:
About Program:
AI/ML for Scientific Discovery Using PyTorch and JAX is a professional training program designed to introduce learners to the use of artificial intelligence and machine learning in modern scientific research. The course focuses on how AI/ML models can support data-driven discovery, pattern recognition, prediction, simulation, and optimization across scientific domains.
Aim: The aim of this course is to equip participants with practical and conceptual knowledge of AI/ML techniques for solving scientific research problems using PyTorch and JAX, enabling them to develop intelligent models for data analysis, prediction, simulation, and discovery.
Program Objectives:
- Understand the role of AI and machine learning in scientific discovery.
- Learn the fundamentals of deep learning workflows using PyTorch and JAX.
- Apply AI/ML models to scientific datasets and research problems.
- Understand how neural networks can be used for prediction, classification, regression, and pattern discovery.
- Explore JAX for high-performance numerical computing and differentiable programming.
- Build basic AI/ML pipelines for scientific data processing and model development.
- Understand the importance of model evaluation, reproducibility, and responsible AI in research.
- Gain exposure to real-world scientific applications of AI/ML across interdisciplinary domains.
What you will learn?
Day 1: Scientific ML Foundations & PyTorch-Based Property Prediction
- AI/ML for scientific discovery, simulation acceleration, and inverse design
- Scientific datasets: molecular, materials, physics, experimental, and simulation data
- Scientific data challenges: small datasets, noise, sparsity, uncertainty, and extrapolation
- Tensor representation of scientific systems using PyTorch
- Automatic differentiation for scientific modeling
- Designing neural networks for scientific property prediction
- Model validation, residual analysis, and scientific interpretability
Hands-on
PyTorch-Based Scientific Property Prediction Model
Participants will build a neural network for predicting a physics, chemistry, or materials-related property using structured scientific data.
Day 2: Physics-Informed Neural Networks & Surrogate Modeling
- Physics-informed neural networks
- Scientific loss functions: data loss, physics loss, constraint loss
- Incorporating differential equations, boundary conditions, and conservation laws
- Surrogate models for expensive simulations
- Parameter estimation and inverse modeling
- Failure modes: non-physical predictions, overfitting, and unstable training
- Applications in diffusion, heat transfer, reaction kinetics, and material response modeling
Hands-on
Physics-Informed Neural Network Prototype
Participants will implement a model that combines data-driven learning with physics-based constraints.
Day 3: JAX for Differentiable Scientific Computing & Optimization
- JAX for high-performance scientific ML
- Key transformations: grad, jit, and vmap
- Differentiable scientific computing and parameter optimization
- Vectorized and accelerated scientific model training
- Optax-based optimization workflows
- Comparing PyTorch and JAX for research applications
- Advanced trends: graph neural networks, neural operators, equivariant models, foundation models, and AI-driven materials discovery
Hands-on
JAX-Based Differentiable Optimization Workflow
Participants will build a JAX model for fitting a physical response curve, optimizing a scientific parameter, or modeling an energy/function landscape.
Mentor Profile
Fee Plan
Get an e-Certificate of Participation!
Intended For :
- Students pursuing science, engineering, computer science, biotechnology, physics, chemistry, mathematics, or related fields
- Ph.D. scholars and research scholars working with scientific or computational data
- Faculty members and academicians interested in AI-driven research
- Industry professionals working in R&D, data science, scientific computing, or applied AI
- Researchers interested in applying PyTorch, JAX, and machine learning to scientific discovery
Career Supporting Skills
Program Outcomes
- Explain how AI/ML accelerates scientific research and innovation.
- Use PyTorch to design, train, and evaluate machine learning/deep learning models.
- Use JAX for numerical computing, automatic differentiation, and scientific machine learning workflows.
- Prepare and process scientific datasets for AI/ML-based analysis.
- Develop basic neural network models for scientific prediction and discovery tasks.
- Interpret model outputs and evaluate model performance using appropriate metrics.
- Apply AI/ML approaches to interdisciplinary research problems.
- Build confidence in using modern AI frameworks for research, experimentation, and scientific computing.