Commits

Branch selector

 main

User selector

All users

All time

Commit History

Commits on Mar 30, 2026

• Bump version to 0.2.

  Show description for 08cf736
  alvarosg

  committed
  08cf736

  Copy full SHA for 08cf736

• Adding xarray_jax vmap.

  Show description for 5dfae39
  alvarosg

  committed
  5dfae39

  Copy full SHA for 5dfae39

• Automated Code Change

  Show description for d39de19
  

  QEDady

  authored and

  alvarosg

  committed
  d39de19

  Copy full SHA for d39de19

• [graphcast] Migrate to the new `jax.pmap`.

  Show description for 12e55eb
  

  danielsuo

  authored and

  alvarosg

  committed
  12e55eb

  Copy full SHA for 12e55eb

• Ensure correct sharding of rng during rollout.

  Show description for 19c8005
  alvarosg

  committed
  19c8005

  Copy full SHA for 19c8005

• Replace unicode escaped characters in ipynb files

  Show description for 1c1a0a6
  

  AleksMat

  authored and

  alvarosg

  committed
  1c1a0a6

  Copy full SHA for 1c1a0a6

• Fix feature utils for update to scipy 1.17

  Show description for d08e86c
  alvarosg

  committed
  d08e86c

  Copy full SHA for d08e86c

• [graphcast] Prepare for `jax_pmap_shmap_merge=True`.

  Show description for 20483b2
  

  danielsuo

  authored and

  alvarosg

  committed
  20483b2

  Copy full SHA for 20483b2

• [pmap] Remove `jax_pmap_shmap_merge` configuration.

  Show description for 40f5d63
  

  danielsuo

  authored and

  alvarosg

  committed
  40f5d63

  Copy full SHA for 40f5d63

Commits on Nov 10, 2025

• Update contact information with privacy policy.

  Show description for c037656
  andrewlkd

  committed
  c037656

  Copy full SHA for c037656

• Update test for `jax_pmap_shmap_merge=True`.

  Show description for 62233a9
  

  danielsuo

  authored and

  andrewlkd

  committed
  62233a9

  Copy full SHA for 62233a9

• Replace `jax.stages.OutInfo` with `jax.ShapeDtypeStruct`

  Show description for 97b9223
  

  yashk2810

  authored and

  andrewlkd

  committed
  97b9223

  Copy full SHA for 97b9223

Commits on Jun 23, 2025

• Clean up math in `get_rotation_matrices_to_local_coordinates`.

  Show description for 5e63ba0
  

  alvarosg

  authored and

  mjwillson

  committed
  5e63ba0

  Copy full SHA for 5e63ba0

• Allowing any leading dims in einsum.

  Show description for 7f6b738
  

  alvarosg

  authored and

  mjwillson

  committed
  7f6b738

  Copy full SHA for 7f6b738

• Explicitly inject numpy and scipy.spatial.transform in some model utils.

  Show description for df43320
  

  alvarosg

  authored and

  mjwillson

  committed
  df43320

  Copy full SHA for df43320

• Adding an `edge_normalization_factor` argument and utils for converting between lat-lon and cartesian.

  Show description for e2ed5fb
  

  alvarosg

  authored and

  mjwillson

  committed
  e2ed5fb

  Copy full SHA for e2ed5fb

• Adds tree_map_with_dims to xarray_jax.

  Show description for 2f7791b
  mjwillson

  committed
  2f7791b

  Copy full SHA for 2f7791b

• Small data_utils.py refactor

  Show description for 4a404e7
  

  alvarosg

  authored and

  mjwillson

  committed
  4a404e7

  Copy full SHA for 4a404e7

• No public description

  Show description for e6d119d
  

  h-joo

  authored and

  mjwillson

  committed
  e6d119d

  Copy full SHA for e6d119d

• Add xarray_jax.scan, which wraps jax.lax.scan to handle xarray inputs/outputs.

  Show description for 9ded766
  mjwillson

  committed
  9ded766

  Copy full SHA for 9ded766

• Update xarray_jax to register xarray.DataTree with jax.tree_util.

  Show description for 677e8d3
  mjwillson

  committed
  677e8d3

  Copy full SHA for 677e8d3

• Avoid call to "_get_next_inputs" in the last step of a rollout.

  Show description for a5f84bb
  

  alvarosg

  authored and

  mjwillson

  committed
  a5f84bb

  Copy full SHA for a5f84bb

• Update xarray_jax to support AOT compilation via `jax.jit(...).lower(...).compile()`, and `jax.eval_shape`.

  Show description for 737b649
  mjwillson

  committed
  737b649

  Copy full SHA for 737b649

Commits on Jan 31, 2025

• Reconfigure JAX in Colab to use latest image.

  Show description for 7077d40
  andrewlkd

  committed
  7077d40

  Copy full SHA for 7077d40

• Ensure sqrt operand does not resolve to negative in apply_stochastic_churn.

  Show description for 4c821a3
  andrewlkd

  committed
  4c821a3

  Copy full SHA for 4c821a3

Commits on Jan 23, 2025

• Use len(pmap_devices) instead of jax.device_count().

  Show description for 49666f4
  

  DeepMind Weather

  authored and

  andrewlkd

  committed
  49666f4

  Copy full SHA for 49666f4

Commits on Dec 17, 2024

• Correct Operational GenCast name in README and some minor formatting.

  andrewlkd

  committed
  6819a0f

  Copy full SHA for 6819a0f

Commits on Dec 13, 2024

• Update cloud_vm_setup with fuller explanations of TPU vs GPU comparisons.

  andrewlkd

  committed
  273af28

  Copy full SHA for 273af28

• Update gencast mini demo notebook defaults.

  andrewlkd

  committed
  10fa386

  Copy full SHA for 10fa386

Commits on Dec 10, 2024

• Move function pmap to separate cell in demo notebooks.

  andrewlkd

  committed
  f71882b

  Copy full SHA for f71882b

• Update GPU documentation and unfreeze SparseTransformerConfig to enable simpler adaptations.

  andrewlkd

  committed
  d1f0a28

  Copy full SHA for d1f0a28

Commits on Dec 4, 2024

• Fix scorecard URLs.

  andrewlkd

  authored
  7855f29

  Copy full SHA for 7855f29

• Move docs to outer repo folder.

  andrewlkd

  committed
  201f2a1

  Copy full SHA for 201f2a1

• Adding GenCast support.

  andrewlkd

  committed
  bf9785f

  Copy full SHA for bf9785f

Commits on Oct 9, 2024

• Bump version to 0.1.1

  alvarosg

  authored
  97d1ad5

  Copy full SHA for 97d1ad5