A highly readable spectral finite element atmospheric dynamical core with automatic differentiation using JAX and/or PyTorch

Autoresearch for weather dycores

Unitful Quantities in JAX

Continuous background code review database for agents, work faster and smarter with accountability for every line of generated code.

Differentiable Fluid Dynamics Package

Experimental Typing Stubs for NumPy

Worktrunk is a CLI for Git worktree management, designed for parallel AI agent workflows

JAX + Units

Coordinate Axes for JAX

JAX compatible datetime and timedelta types

Open-source, cloud-native transactional tensor storage engine

Ai2 Climate Emulator: fast machine learning models for weather and climate prediction

Jatmos is a scalable simulator for atmospheric dynamics, written in JAX.

Minimal JAX implementation of k-nearest neighbors using a k-d tree.

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

Universal Notation for Tensor Operations in 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/