Welcome to Jeo¶

Jax library to support machine learning research for remote sensing and Earth observation.

Overview¶

Jeo is a specialized open source framework developed by Google DeepMind that accelerates machine learning for geospatial remote sensing and earth observation (EO) tasks. It uses JAX and Flax for high-performance model training on large geospatial datasets.

Jeo is tailored to the characteristics of geospatial datasets to help in the development of models that can operate at scale. Jeo is primarily intended for researchers and developers actively working in fields such as geospatial analysis, remote sensing, environmental science, and sustainability modelling.

To access Jeo, visit the GitHub repository.

Design¶

Jeo code structure is inspired by Big Vision and Scenic, and it builds upon the following:

JAX: provides the engine for high-performance computation.
Flax: offers tools for building neural network models.
tf.data: manages data input pipelines.
GeeFlow: connects the framework to Google Earth Engine's data resources.

Familiarity with these components can aid in customizing and extending Jeo for specific research needs. By using JAX and Flax, Jeo inherently benefits from features like automatic differentiation, code compilation (XLA), and seamless execution across various hardware accelerators, including 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.

Furthermore, Jeo's effective integration with libraries like GeeFlow enables efficient construction of large-scale datasets directly from Google Earth Engine (GEE), streamlining workflows in the GEE ecosystem. This combination makes Jeo a potent tool for researchers pushing the boundaries of AI applications in understanding and modelling our planet for research and sustainability projects.

Cite Jeo¶

Cite the Jeo codebase as follows:

@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}
}