Welcome to Jeo¶

Geospatial ML — from satellite imagery to trained models.

What is Jeo?¶

Jeo is an open source framework developed by Google DeepMind for machine learning on geospatial remote sensing and Earth observation (EO) data. It uses JAX and Flax for high-performance model training and inference on large-scale geospatial datasets.

Jeo is primarily intended for researchers and developers working in geospatial analysis, remote sensing, environmental science, and sustainability modelling.

Where to Start¶

Your Goal	Go To
🌍 New to Jeo? Complete end-to-end tutorial	Starter Tutorial
⚙️ Understand configs and how to write them	Configs
🚀 Launch an experiment on XManager or locally	Launching
📦 Start a new project from scratch	New Project Guide
🔧 Troubleshooting common errors	FAQ

Documentation¶

Getting Started¶

Starter Tutorial — A hands-on, phase-by-phase walkthrough using a real dataset (GlobalGeoTree). Covers the full pipeline from data discovery through GeeFlow export, model training, and evaluation.

Core Concepts¶

Configs — Config system, custom modules, duration specs, sweeps, Vizier integration, and loading pretrained models.
Models — Model architectures, input/output conventions, weight loading, and module resolution.
Tasks & Losses — TaskBase interface, loss functions, and available task types.
Datasets — TFDS datasets, custom loaders, and GeeFlow integration.
Preprocessing — The pipe-based preprocessing grammar and available operations.
Evaluators — Evaluator interface, metrics, Lumascope visualization.

Guides¶

Reference¶

Conventions & Style — Coding conventions, abbreviations, config style guide, and design principles.
FAQ & Troubleshooting — Common errors, memory limits, GPU OOM, GeeFlow authentication, and debugging tips.

Design¶

Jeo's code structure is inspired by Big Vision and Scenic, and builds upon:

JAX — High-performance computation engine.
Flax — Neural network building tools.
tf.data — Data input pipelines.
GeeFlow — Connects to Google Earth Engine's satellite data resources for large-scale dataset generation.

By using JAX and Flax, Jeo benefits from automatic differentiation, XLA compilation, and seamless execution across CPUs, GPUs, and Google Cloud TPUs. This focus on performance is particularly relevant for Earth observation tasks, which often involve processing massive datasets derived from satellite imagery and other remote sensing platforms.

Cite Jeo¶

@software{jeo2025:github,
  author = {Maxim Neumann and Anton Raichuk and Michelangelo Conserva and
  Luis Miguel Pazos-Outón and Keith Anderson and Matt Overlan and Mélanie Rey
  and Yuchang Jiang and Petra Poklukar and Cristina Nader Vasconcelos},
  title = {{JEO}: Model training and inference for geospatial remote sensing and
  {E}arth {O}bservation in {JAX}},
  url = {https://github.com/google-deepmind/jeo},
  year = {2025}
}