Open-source, cloud-native transactional tensor storage engine

scores: Metrics for the verification, evaluation and optimisation of forecasts, predictions or models.

Universal Tensor Operations in Einstein-Inspired Notation for Python.

The versatile ocean simulator, in pure Python, powered by JAX.

NeuralGCM's differentiable dycore

An Xarray extension for Google Earth Engine

A benchmark for the next generation of data-driven global weather models.

MVP for opening xarray datasets from cloud zarr stores in a single thread with asyncio loops

Neural Networks for JAX

A dead simple Python package for creating custom JAX pytree objects

Stencil computations in JAX

JAX-based neural network library

Elegant easy-to-use neural networks + scientific computing in JAX. https://docs.kidger.site/equinox/

jax library for E3 Equivariant Neural Networks

Scalable array processing with bounded memory

A benchmark dataset for data-driven weather forecasting

⚡ xarray data creation by data classes

Numerical differential equation solvers in JAX. Autodifferentiable and GPU-capable. https://docs.kidger.site/diffrax/

Distributed Xarray with Apache Beam

PRQL is a modern language for transforming data — a simple, powerful, pipelined SQL replacement

A debugging and profiling tool that can trace and visualize python code execution

Play atmospheric modelling like it's LEGO.

Mathematical operations for JAX pytrees

Tools to make weather data accessible and useful.

Zero-copy MPI communication of JAX arrays, for turbo-charged HPC applications in Python ⚡

Fast & furious GroupBy operations for dask.array

Computational Fluid Dynamics in JAX

A differentiable PDE solving framework for machine learning