From 7453ae144bbd098ff14f6b25ba6d7de26ba9efca Mon Sep 17 00:00:00 2001
From: "William S. Moses" <gh@wsmoses.com>
Date: Mon, 14 Oct 2024 13:08:59 -0500
Subject: [PATCH 01/17] add neuralgcm

---
 builddeps/test-requirements.txt |   1 +
 test/BUILD                      |  12 +++
 test/jaxmd.py                   | 109 +-------------------
 test/llama.py                   | 106 --------------------
 test/maxtext.py                 | 105 --------------------
 test/neuralgcm.py               | 169 ++++++++++++++++++++++++++++++++
 test/test_utils.py              | 106 ++++++++++++++++++++
 7 files changed, 293 insertions(+), 315 deletions(-)
 create mode 100644 test/neuralgcm.py

diff --git a/builddeps/test-requirements.txt b/builddeps/test-requirements.txt
index 159b9152..69e723fa 100644
--- a/builddeps/test-requirements.txt
+++ b/builddeps/test-requirements.txt
@@ -10,6 +10,7 @@ https://github.com/wsmoses/jraph/archive/b00d9a03db76c69a258a86df81638b9a2f28829
 # https://github.com/wsmoses/maxtext/archive/bc50722be7d89e4003bd830b80e4ac968be658eb.tar.gz
 
 jax-cuda12-plugin[with_cuda]; sys_platform == 'linux'
+neuralgcm
 requests; sys_platform == 'linux'
 
 # -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
diff --git a/test/BUILD b/test/BUILD
index c57ef6df..4ae06796 100644
--- a/test/BUILD
+++ b/test/BUILD
@@ -136,6 +136,18 @@ py_test(
     timeout='long'
 )
 
+py_test(
+    name = "neuralgcm",
+    srcs = [
+        "neuralgcm.py",
+        "test_utils.py",
+    ],
+    imports = ["."],
+    deps = TEST_DEPS + ["@pypi_neuralgcm//:pkg"],
+    timeout='long'
+)
+
+
 # py_test(
 #     name = "maxtext",
 #     srcs = [
diff --git a/test/jaxmd.py b/test/jaxmd.py
index b360100e..0a4a22bc 100644
--- a/test/jaxmd.py
+++ b/test/jaxmd.py
@@ -24,110 +24,11 @@
 
 jax.config.update("jax_enable_x64", True)
 
-partialopt = (
-    "inline{default-pipeline=canonicalize max-iterations=4},"
-    + """canonicalize,cse,
-enzyme-hlo-generate-td{
-            patterns=compare_op_canon<16>;
-transpose_transpose<16>;
-broadcast_in_dim_op_canon<16>;
-convert_op_canon<16>;
-dynamic_broadcast_in_dim_op_not_actually_dynamic<16>;
-chained_dynamic_broadcast_in_dim_canonicalization<16>;
-dynamic_broadcast_in_dim_all_dims_non_expanding<16>;
-noop_reduce_op_canon<16>;
-empty_reduce_op_canon<16>;
-dynamic_reshape_op_canon<16>;
-get_tuple_element_op_canon<16>;
-real_op_canon<16>;
-imag_op_canon<16>;
-get_dimension_size_op_canon<16>;
-gather_op_canon<16>;
-reshape_op_canon<16>;
-merge_consecutive_reshapes<16>;
-transpose_is_reshape<16>;
-zero_extent_tensor_canon<16>;
-reorder_elementwise_and_shape_op<16>;
-
-cse_broadcast_in_dim<16>;
-cse_slice<16>;
-cse_transpose<16>;
-cse_convert<16>;
-cse_pad<16>;
-cse_dot_general<16>;
-cse_reshape<16>;
-cse_mul<16>;
-cse_div<16>;
-cse_add<16>;
-cse_subtract<16>;
-cse_min<16>;
-cse_max<16>;
-cse_neg<16>;
-cse_concatenate<16>;
-
-concatenate_op_canon<16>(1024);
-select_op_canon<16>(1024);
-add_simplify<16>;
-sub_simplify<16>;
-and_simplify<16>;
-max_simplify<16>;
-min_simplify<16>;
-or_simplify<16>;
-negate_simplify<16>;
-mul_simplify<16>;
-div_simplify<16>;
-rem_simplify<16>;
-pow_simplify<16>;
-sqrt_simplify<16>;
-cos_simplify<16>;
-sin_simplify<16>;
-noop_slice<16>;
-const_prop_through_barrier<16>;
-slice_slice<16>;
-shift_right_logical_simplify<16>;
-pad_simplify<16>;
-negative_pad_to_slice<16>;
-tanh_simplify<16>;
-exp_simplify<16>;
-slice_simplify<16>;
-convert_simplify<16>;
-dynamic_slice_to_static<16>;
-dynamic_update_slice_elim<16>;
-concat_to_broadcast<16>;
-reduce_to_reshape<16>;
-broadcast_to_reshape<16>;
-gather_simplify<16>;
-iota_simplify<16>(1024);
-broadcast_in_dim_simplify<16>(1024);
-convert_concat<1>;
-dynamic_update_to_concat<1>;
-slice_of_dynamic_update<1>;
-slice_elementwise<1>;
-slice_pad<1>;
-dot_reshape_dot<1>;
-concat_const_prop<1>;
-concat_fuse<1>;
-pad_reshape_pad<1>;
-pad_pad<1>;
-concat_push_binop_add<1>;
-concat_push_binop_mul<1>;
-scatter_to_dynamic_update_slice<1>;
-reduce_concat<1>;
-slice_concat<1>;
-
-bin_broadcast_splat_add<1>;
-bin_broadcast_splat_subtract<1>;
-bin_broadcast_splat_div<1>;
-bin_broadcast_splat_mul<1>;
-slice_reshape<1>;
-
-dot_reshape_pad<1>;
-pad_dot_general<1>(1);
-broadcast_reduce<1>;
-            },
-            transform-interpreter,
-            enzyme-hlo-remove-transform,cse"""
-)
+from jax_md import space
+from jax_md import energy
+from jax_md import simulate
+from jax_md import quantity
+from jax_md import partition
 
 pipelines = [
     ("JaX  ", None, CurBackends),
diff --git a/test/llama.py b/test/llama.py
index 6effbe46..84b1fdf4 100644
--- a/test/llama.py
+++ b/test/llama.py
@@ -247,112 +247,6 @@ def forward(x, config, weights, key_cache, value_cache):
 
     return x
 
-
-partialopt = (
-    "inline{default-pipeline=canonicalize max-iterations=4},"
-    + """canonicalize,cse,
-enzyme-hlo-generate-td{
-            patterns=compare_op_canon<16>;
-transpose_transpose<16>;
-broadcast_in_dim_op_canon<16>;
-convert_op_canon<16>;
-dynamic_broadcast_in_dim_op_not_actually_dynamic<16>;
-chained_dynamic_broadcast_in_dim_canonicalization<16>;
-dynamic_broadcast_in_dim_all_dims_non_expanding<16>;
-noop_reduce_op_canon<16>;
-empty_reduce_op_canon<16>;
-dynamic_reshape_op_canon<16>;
-get_tuple_element_op_canon<16>;
-real_op_canon<16>;
-imag_op_canon<16>;
-get_dimension_size_op_canon<16>;
-gather_op_canon<16>;
-reshape_op_canon<16>;
-merge_consecutive_reshapes<16>;
-transpose_is_reshape<16>;
-zero_extent_tensor_canon<16>;
-reorder_elementwise_and_shape_op<16>;
-
-cse_broadcast_in_dim<16>;
-cse_slice<16>;
-cse_transpose<16>;
-cse_convert<16>;
-cse_pad<16>;
-cse_dot_general<16>;
-cse_reshape<16>;
-cse_mul<16>;
-cse_div<16>;
-cse_add<16>;
-cse_subtract<16>;
-cse_min<16>;
-cse_max<16>;
-cse_neg<16>;
-cse_concatenate<16>;
-
-concatenate_op_canon<16>(1024);
-select_op_canon<16>(1024);
-add_simplify<16>;
-sub_simplify<16>;
-and_simplify<16>;
-max_simplify<16>;
-min_simplify<16>;
-or_simplify<16>;
-negate_simplify<16>;
-mul_simplify<16>;
-div_simplify<16>;
-rem_simplify<16>;
-pow_simplify<16>;
-sqrt_simplify<16>;
-cos_simplify<16>;
-sin_simplify<16>;
-noop_slice<16>;
-const_prop_through_barrier<16>;
-slice_slice<16>;
-shift_right_logical_simplify<16>;
-pad_simplify<16>;
-negative_pad_to_slice<16>;
-tanh_simplify<16>;
-exp_simplify<16>;
-slice_simplify<16>;
-convert_simplify<16>;
-dynamic_slice_to_static<16>;
-dynamic_update_slice_elim<16>;
-concat_to_broadcast<16>;
-reduce_to_reshape<16>;
-broadcast_to_reshape<16>;
-gather_simplify<16>;
-iota_simplify<16>(1024);
-broadcast_in_dim_simplify<16>(1024);
-convert_concat<1>;
-dynamic_update_to_concat<1>;
-slice_of_dynamic_update<1>;
-slice_elementwise<1>;
-slice_pad<1>;
-dot_reshape_dot<1>;
-concat_const_prop<1>;
-concat_fuse<1>;
-pad_reshape_pad<1>;
-pad_pad<1>;
-concat_push_binop_add<1>;
-concat_push_binop_mul<1>;
-scatter_to_dynamic_update_slice<1>;
-reduce_concat<1>;
-slice_concat<1>;
-
-bin_broadcast_splat_add<1>;
-bin_broadcast_splat_subtract<1>;
-bin_broadcast_splat_div<1>;
-bin_broadcast_splat_mul<1>;
-slice_reshape<1>;
-
-dot_reshape_pad<1>;
-pad_dot_general<1>(1);
-broadcast_reduce<1>;
-            },
-            transform-interpreter,
-            enzyme-hlo-remove-transform,cse"""
-)
-
 pipelines = [
     ("JaX  ", None, CurBackends),
     ("JaXPipe", JaXPipeline(), CurBackends),
diff --git a/test/maxtext.py b/test/maxtext.py
index b14c1874..2f2a9267 100644
--- a/test/maxtext.py
+++ b/test/maxtext.py
@@ -28,111 +28,6 @@
 from jax import random
 from jax import lax
 
-partialopt = (
-    "inline{default-pipeline=canonicalize max-iterations=4},"
-    + """canonicalize,cse,
-enzyme-hlo-generate-td{
-            patterns=compare_op_canon<16>;
-transpose_transpose<16>;
-broadcast_in_dim_op_canon<16>;
-convert_op_canon<16>;
-dynamic_broadcast_in_dim_op_not_actually_dynamic<16>;
-chained_dynamic_broadcast_in_dim_canonicalization<16>;
-dynamic_broadcast_in_dim_all_dims_non_expanding<16>;
-noop_reduce_op_canon<16>;
-empty_reduce_op_canon<16>;
-dynamic_reshape_op_canon<16>;
-get_tuple_element_op_canon<16>;
-real_op_canon<16>;
-imag_op_canon<16>;
-get_dimension_size_op_canon<16>;
-gather_op_canon<16>;
-reshape_op_canon<16>;
-merge_consecutive_reshapes<16>;
-transpose_is_reshape<16>;
-zero_extent_tensor_canon<16>;
-reorder_elementwise_and_shape_op<16>;
-
-cse_broadcast_in_dim<16>;
-cse_slice<16>;
-cse_transpose<16>;
-cse_convert<16>;
-cse_pad<16>;
-cse_dot_general<16>;
-cse_reshape<16>;
-cse_mul<16>;
-cse_div<16>;
-cse_add<16>;
-cse_subtract<16>;
-cse_min<16>;
-cse_max<16>;
-cse_neg<16>;
-cse_concatenate<16>;
-
-concatenate_op_canon<16>(1024);
-select_op_canon<16>(1024);
-add_simplify<16>;
-sub_simplify<16>;
-and_simplify<16>;
-max_simplify<16>;
-min_simplify<16>;
-or_simplify<16>;
-negate_simplify<16>;
-mul_simplify<16>;
-div_simplify<16>;
-rem_simplify<16>;
-pow_simplify<16>;
-sqrt_simplify<16>;
-cos_simplify<16>;
-sin_simplify<16>;
-noop_slice<16>;
-const_prop_through_barrier<16>;
-slice_slice<16>;
-shift_right_logical_simplify<16>;
-pad_simplify<16>;
-negative_pad_to_slice<16>;
-tanh_simplify<16>;
-exp_simplify<16>;
-slice_simplify<16>;
-convert_simplify<16>;
-dynamic_slice_to_static<16>;
-dynamic_update_slice_elim<16>;
-concat_to_broadcast<16>;
-reduce_to_reshape<16>;
-broadcast_to_reshape<16>;
-gather_simplify<16>;
-iota_simplify<16>(1024);
-broadcast_in_dim_simplify<16>(1024);
-convert_concat<1>;
-dynamic_update_to_concat<1>;
-slice_of_dynamic_update<1>;
-slice_elementwise<1>;
-slice_pad<1>;
-dot_reshape_dot<1>;
-concat_const_prop<1>;
-concat_fuse<1>;
-pad_reshape_pad<1>;
-pad_pad<1>;
-concat_push_binop_add<1>;
-concat_push_binop_mul<1>;
-scatter_to_dynamic_update_slice<1>;
-reduce_concat<1>;
-slice_concat<1>;
-
-bin_broadcast_splat_add<1>;
-bin_broadcast_splat_subtract<1>;
-bin_broadcast_splat_div<1>;
-bin_broadcast_splat_mul<1>;
-slice_reshape<1>;
-
-dot_reshape_pad<1>;
-pad_dot_general<1>(1);
-broadcast_reduce<1>;
-            },
-            transform-interpreter,
-            enzyme-hlo-remove-transform,cse"""
-)
-
 pipelines = [
     ("JaX  ", None),
     ("JaXPipe", JaXPipeline()),
diff --git a/test/neuralgcm.py b/test/neuralgcm.py
new file mode 100644
index 00000000..14049505
--- /dev/null
+++ b/test/neuralgcm.py
@@ -0,0 +1,169 @@
+from absl.testing import absltest
+import jax.numpy as jnp
+import jax.random
+import jax.lax
+import enzyme_ad.jax as enzyme_jax
+from enzyme_ad.jax import (
+    enzyme_jax_ir,
+    NewXLAPipeline,
+    OldXLAPipeline,
+    JaXPipeline,
+    hlo_opts,
+)
+import numpy as np
+import timeit
+from test_utils import *
+
+argv = ("-I/usr/include/c++/11", "-I/usr/include/x86_64-linux-gnu/c++/11")
+
+import jax.numpy as np
+import numpy as onp
+from jax import jit
+from jax import random
+from jax import lax
+
+jax.config.update("jax_enable_x64", True)
+
+from jax_md import space
+from jax_md import energy
+from jax_md import simulate
+from jax_md import quantity
+from jax_md import partition
+
+pipelines = [
+    ("JaX  ", None, CurBackends),
+    ("JaXPipe", JaXPipeline(), CurBackends),
+    (
+        "HLOOpt",
+        JaXPipeline(
+            "inline{default-pipeline=canonicalize max-iterations=4},"
+            + "canonicalize,cse,enzyme-hlo-opt,cse"
+        ),
+        CurBackends,
+    ),
+    ("PartOpt", JaXPipeline(partialopt), CurBackends),
+    ("DefOpt", JaXPipeline(hlo_opts()), CurBackends),
+]
+
+import gcsfs
+import jax
+import numpy as np
+import pickle
+import xarray
+import timeit
+
+from dinosaur import horizontal_interpolation
+from dinosaur import spherical_harmonic
+from dinosaur import xarray_utils
+import neuralgcm
+
+gcs = gcsfs.GCSFileSystem(token='anon')
+
+model_name = 'neural_gcm_dynamic_forcing_deterministic_1_4_deg.pkl'  #@param ['neural_gcm_dynamic_forcing_deterministic_0_7_deg.pkl', 'neural_gcm_dynamic_forcing_deterministic_1_4_deg.pkl', 'neural_gcm_dynamic_forcing_deterministic_2_8_deg.pkl', 'neural_gcm_dynamic_forcing_stochastic_1_4_deg.pkl'] {type: "string"}
+
+with gcs.open(f'gs://gresearch/neuralgcm/04_30_2024/{model_name}', 'rb') as f:
+  ckpt = pickle.load(f)
+
+model = neuralgcm.PressureLevelModel.from_checkpoint(ckpt)
+
+era5_path = 'gs://gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3'
+full_era5 = xarray.open_zarr(gcs.get_mapper(era5_path), chunks=None)
+
+demo_start_time = '2020-02-14'
+demo_end_time = '2020-02-18'
+data_inner_steps = 24  # process every 24th hour
+
+sliced_era5 = (
+    full_era5
+    [model.input_variables + model.forcing_variables]
+    .pipe(
+        xarray_utils.selective_temporal_shift,
+        variables=model.forcing_variables,
+        time_shift='24 hours',
+    )
+    .sel(time=slice(demo_start_time, demo_end_time, data_inner_steps))
+    .compute()
+)
+
+era5_grid = spherical_harmonic.Grid(
+    latitude_nodes=full_era5.sizes['latitude'],
+    longitude_nodes=full_era5.sizes['longitude'],
+    latitude_spacing=xarray_utils.infer_latitude_spacing(full_era5.latitude),
+    longitude_offset=xarray_utils.infer_longitude_offset(full_era5.longitude),
+)
+regridder = horizontal_interpolation.ConservativeRegridder(
+    era5_grid, model.data_coords.horizontal, skipna=True
+)
+eval_era5 = xarray_utils.regrid(sliced_era5, regridder)
+eval_era5 = xarray_utils.fill_nan_with_nearest(eval_era5)
+
+inner_steps = 24  # save model outputs once every 24 hours
+outer_steps = 4 * 24 // inner_steps  # total of 4 days
+timedelta = np.timedelta64(1, 'h') * inner_steps
+times = (np.arange(outer_steps) * inner_steps)  # time axis in hours
+
+# initialize model state
+inputs = model.inputs_from_xarray(eval_era5.isel(time=0))
+input_forcings = model.forcings_from_xarray(eval_era5.isel(time=0))
+rng_key = jax.random.key(42)  # optional for deterministic models
+initial_state = model.encode(inputs, input_forcings, rng_key)
+
+# use persistence for forcing variables (SST and sea ice cover)
+all_forcings = model.forcings_from_xarray(eval_era5.head(time=1))
+
+# make forecast
+final_state, predictions = model.unroll(
+    initial_state,
+    all_forcings,
+    steps=outer_steps,
+    timedelta=timedelta,
+    start_with_input=True,
+)
+predictions_ds = model.data_to_xarray(predictions, times=times)
+
+
+import enzyme_ad.jax as enzyme_jax
+
+from enzyme_ad.jax import (
+    enzyme_jax_ir,
+    NewXLAPipeline,
+    OldXLAPipeline,
+    JaXPipeline,
+    hlo_opts,
+)
+
+def sub(initial_state, all_forcings):
+    return model.unroll(initial_state,
+    all_forcings,
+    steps=outer_steps,
+    timedelta=timedelta,
+    start_with_input=True,
+    )
+
+def run(fn, steps=1):
+    inputs = model.inputs_from_xarray(eval_era5.isel(time=0))
+    input_forcings = model.forcings_from_xarray(eval_era5.isel(time=0))
+    rng_key = jax.random.key(42)  # optional for deterministic models
+    initial_state = model.encode(inputs, input_forcings, rng_key)
+    timedelta = np.timedelta64(1, 'h') * inner_steps
+    map(lambda x:x.block_until_ready(), fn(
+        initial_state,
+        all_forcings,
+    ))
+    return timeit.Timer(
+"""map(lambda x:x.block_until_ready(), fn(
+    initial_state,
+    all_forcings,
+))""",
+        globals={
+            "fn": fn,
+            "initial_state":initial_state,
+            "timedelta":timedelta,
+            "all_forcings":all_forcings,
+            "outer_steps":outer_steps
+        }
+    ).timeit(steps)
+
+run(jax.jit(sub))
+
+run(jax.jit(enzyme_jax_ir(pipeline_options=JaXPipeline(hlo_opts()))(sub)))
diff --git a/test/test_utils.py b/test/test_utils.py
index 78cf634a..c95f6f96 100644
--- a/test/test_utils.py
+++ b/test/test_utils.py
@@ -39,6 +39,112 @@ def setup_backends():
     ("OldXLA", OldXLAPipeline(), ["cpu"]),
 ]
 
+partialopt = (
+    "inline{default-pipeline=canonicalize max-iterations=4},"
+    + """canonicalize,cse,
+enzyme-hlo-generate-td{
+            patterns=compare_op_canon<16>;
+transpose_transpose<16>;
+broadcast_in_dim_op_canon<16>;
+convert_op_canon<16>;
+dynamic_broadcast_in_dim_op_not_actually_dynamic<16>;
+chained_dynamic_broadcast_in_dim_canonicalization<16>;
+dynamic_broadcast_in_dim_all_dims_non_expanding<16>;
+noop_reduce_op_canon<16>;
+empty_reduce_op_canon<16>;
+dynamic_reshape_op_canon<16>;
+get_tuple_element_op_canon<16>;
+real_op_canon<16>;
+imag_op_canon<16>;
+get_dimension_size_op_canon<16>;
+gather_op_canon<16>;
+reshape_op_canon<16>;
+merge_consecutive_reshapes<16>;
+transpose_is_reshape<16>;
+zero_extent_tensor_canon<16>;
+reorder_elementwise_and_shape_op<16>;
+
+cse_broadcast_in_dim<16>;
+cse_slice<16>;
+cse_transpose<16>;
+cse_convert<16>;
+cse_pad<16>;
+cse_dot_general<16>;
+cse_reshape<16>;
+cse_mul<16>;
+cse_div<16>;
+cse_add<16>;
+cse_subtract<16>;
+cse_min<16>;
+cse_max<16>;
+cse_neg<16>;
+cse_concatenate<16>;
+
+concatenate_op_canon<16>(1024);
+select_op_canon<16>(1024);
+add_simplify<16>;
+sub_simplify<16>;
+and_simplify<16>;
+max_simplify<16>;
+min_simplify<16>;
+or_simplify<16>;
+negate_simplify<16>;
+mul_simplify<16>;
+div_simplify<16>;
+rem_simplify<16>;
+pow_simplify<16>;
+sqrt_simplify<16>;
+cos_simplify<16>;
+sin_simplify<16>;
+noop_slice<16>;
+const_prop_through_barrier<16>;
+slice_slice<16>;
+shift_right_logical_simplify<16>;
+pad_simplify<16>;
+negative_pad_to_slice<16>;
+tanh_simplify<16>;
+exp_simplify<16>;
+slice_simplify<16>;
+convert_simplify<16>;
+dynamic_slice_to_static<16>;
+dynamic_update_slice_elim<16>;
+concat_to_broadcast<16>;
+reduce_to_reshape<16>;
+broadcast_to_reshape<16>;
+gather_simplify<16>;
+iota_simplify<16>(1024);
+broadcast_in_dim_simplify<16>(1024);
+convert_concat<1>;
+dynamic_update_to_concat<1>;
+slice_of_dynamic_update<1>;
+slice_elementwise<1>;
+slice_pad<1>;
+dot_reshape_dot<1>;
+concat_const_prop<1>;
+concat_fuse<1>;
+pad_reshape_pad<1>;
+pad_pad<1>;
+concat_push_binop_add<1>;
+concat_push_binop_mul<1>;
+scatter_to_dynamic_update_slice<1>;
+reduce_concat<1>;
+slice_concat<1>;
+
+bin_broadcast_splat_add<1>;
+bin_broadcast_splat_subtract<1>;
+bin_broadcast_splat_div<1>;
+bin_broadcast_splat_mul<1>;
+slice_reshape<1>;
+
+dot_reshape_pad<1>;
+pad_dot_general<1>(1);
+broadcast_reduce<1>;
+            },
+            transform-interpreter,
+            enzyme-hlo-remove-transform,cse"""
+)
+
+
 
 def no_newxla(x):
     return [

From c75f7e7347ea5049774c20013b3a5d9d1558f53c Mon Sep 17 00:00:00 2001
From: "William S. Moses" <gh@wsmoses.com>
Date: Wed, 16 Oct 2024 00:30:45 -0400
Subject: [PATCH 02/17] fix

---
 builddeps/test-requirements.txt |   1 +
 test/BUILD                      |   2 +-
 test/neuralgcm.py               | 223 ++++++++++++++++----------------
 3 files changed, 110 insertions(+), 116 deletions(-)

diff --git a/builddeps/test-requirements.txt b/builddeps/test-requirements.txt
index 69e723fa..4b99f73e 100644
--- a/builddeps/test-requirements.txt
+++ b/builddeps/test-requirements.txt
@@ -11,6 +11,7 @@ https://github.com/wsmoses/jraph/archive/b00d9a03db76c69a258a86df81638b9a2f28829
 
 jax-cuda12-plugin[with_cuda]; sys_platform == 'linux'
 neuralgcm
+gcsfs
 requests; sys_platform == 'linux'
 
 # -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
diff --git a/test/BUILD b/test/BUILD
index 4ae06796..9fda413f 100644
--- a/test/BUILD
+++ b/test/BUILD
@@ -143,7 +143,7 @@ py_test(
         "test_utils.py",
     ],
     imports = ["."],
-    deps = TEST_DEPS + ["@pypi_neuralgcm//:pkg"],
+    deps = TEST_DEPS + ["@pypi_neuralgcm//:pkg", "@pypi_gcsfs//:pkg"],
     timeout='long'
 )
 
diff --git a/test/neuralgcm.py b/test/neuralgcm.py
index 14049505..e2d58d5e 100644
--- a/test/neuralgcm.py
+++ b/test/neuralgcm.py
@@ -22,13 +22,6 @@
 from jax import random
 from jax import lax
 
-jax.config.update("jax_enable_x64", True)
-
-from jax_md import space
-from jax_md import energy
-from jax_md import simulate
-from jax_md import quantity
-from jax_md import partition
 
 pipelines = [
     ("JaX  ", None, CurBackends),
@@ -57,113 +50,113 @@
 from dinosaur import xarray_utils
 import neuralgcm
 
-gcs = gcsfs.GCSFileSystem(token='anon')
-
-model_name = 'neural_gcm_dynamic_forcing_deterministic_1_4_deg.pkl'  #@param ['neural_gcm_dynamic_forcing_deterministic_0_7_deg.pkl', 'neural_gcm_dynamic_forcing_deterministic_1_4_deg.pkl', 'neural_gcm_dynamic_forcing_deterministic_2_8_deg.pkl', 'neural_gcm_dynamic_forcing_stochastic_1_4_deg.pkl'] {type: "string"}
-
-with gcs.open(f'gs://gresearch/neuralgcm/04_30_2024/{model_name}', 'rb') as f:
-  ckpt = pickle.load(f)
-
-model = neuralgcm.PressureLevelModel.from_checkpoint(ckpt)
-
-era5_path = 'gs://gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3'
-full_era5 = xarray.open_zarr(gcs.get_mapper(era5_path), chunks=None)
-
-demo_start_time = '2020-02-14'
-demo_end_time = '2020-02-18'
-data_inner_steps = 24  # process every 24th hour
-
-sliced_era5 = (
-    full_era5
-    [model.input_variables + model.forcing_variables]
-    .pipe(
-        xarray_utils.selective_temporal_shift,
-        variables=model.forcing_variables,
-        time_shift='24 hours',
-    )
-    .sel(time=slice(demo_start_time, demo_end_time, data_inner_steps))
-    .compute()
-)
-
-era5_grid = spherical_harmonic.Grid(
-    latitude_nodes=full_era5.sizes['latitude'],
-    longitude_nodes=full_era5.sizes['longitude'],
-    latitude_spacing=xarray_utils.infer_latitude_spacing(full_era5.latitude),
-    longitude_offset=xarray_utils.infer_longitude_offset(full_era5.longitude),
-)
-regridder = horizontal_interpolation.ConservativeRegridder(
-    era5_grid, model.data_coords.horizontal, skipna=True
-)
-eval_era5 = xarray_utils.regrid(sliced_era5, regridder)
-eval_era5 = xarray_utils.fill_nan_with_nearest(eval_era5)
-
-inner_steps = 24  # save model outputs once every 24 hours
-outer_steps = 4 * 24 // inner_steps  # total of 4 days
-timedelta = np.timedelta64(1, 'h') * inner_steps
-times = (np.arange(outer_steps) * inner_steps)  # time axis in hours
-
-# initialize model state
-inputs = model.inputs_from_xarray(eval_era5.isel(time=0))
-input_forcings = model.forcings_from_xarray(eval_era5.isel(time=0))
-rng_key = jax.random.key(42)  # optional for deterministic models
-initial_state = model.encode(inputs, input_forcings, rng_key)
-
-# use persistence for forcing variables (SST and sea ice cover)
-all_forcings = model.forcings_from_xarray(eval_era5.head(time=1))
-
-# make forecast
-final_state, predictions = model.unroll(
-    initial_state,
-    all_forcings,
-    steps=outer_steps,
-    timedelta=timedelta,
-    start_with_input=True,
-)
-predictions_ds = model.data_to_xarray(predictions, times=times)
-
-
-import enzyme_ad.jax as enzyme_jax
-
-from enzyme_ad.jax import (
-    enzyme_jax_ir,
-    NewXLAPipeline,
-    OldXLAPipeline,
-    JaXPipeline,
-    hlo_opts,
-)
-
-def sub(initial_state, all_forcings):
-    return model.unroll(initial_state,
-    all_forcings,
-    steps=outer_steps,
-    timedelta=timedelta,
-    start_with_input=True,
-    )
-
-def run(fn, steps=1):
-    inputs = model.inputs_from_xarray(eval_era5.isel(time=0))
-    input_forcings = model.forcings_from_xarray(eval_era5.isel(time=0))
-    rng_key = jax.random.key(42)  # optional for deterministic models
-    initial_state = model.encode(inputs, input_forcings, rng_key)
-    timedelta = np.timedelta64(1, 'h') * inner_steps
-    map(lambda x:x.block_until_ready(), fn(
+class NeuralGCM(absltest.TestCase):
+    def setUp(self):
+        gcs = gcsfs.GCSFileSystem(token='anon')
+
+        model_name = 'neural_gcm_dynamic_forcing_deterministic_1_4_deg.pkl'  #@param ['neural_gcm_dynamic_forcing_deterministic_0_7_deg.pkl', 'neural_gcm_dynamic_forcing_deterministic_1_4_deg.pkl', 'neural_gcm_dynamic_forcing_deterministic_2_8_deg.pkl', 'neural_gcm_dynamic_forcing_stochastic_1_4_deg.pkl'] {type: "string"}
+
+        with gcs.open(f'gs://gresearch/neuralgcm/04_30_2024/{model_name}', 'rb') as f:
+          ckpt = pickle.load(f)
+
+        model = neuralgcm.PressureLevelModel.from_checkpoint(ckpt)
+
+        era5_path = 'gs://gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3'
+        full_era5 = xarray.open_zarr(gcs.get_mapper(era5_path), chunks=None)
+
+        demo_start_time = '2020-02-14'
+        demo_end_time = '2020-02-18'
+        data_inner_steps = 24  # process every 24th hour
+
+        sliced_era5 = (
+            full_era5
+            [model.input_variables + model.forcing_variables]
+            .pipe(
+                xarray_utils.selective_temporal_shift,
+                variables=model.forcing_variables,
+                time_shift='24 hours',
+            )
+            .sel(time=slice(demo_start_time, demo_end_time, data_inner_steps))
+            .compute()
+        )
+
+        era5_grid = spherical_harmonic.Grid(
+            latitude_nodes=full_era5.sizes['latitude'],
+            longitude_nodes=full_era5.sizes['longitude'],
+            latitude_spacing=xarray_utils.infer_latitude_spacing(full_era5.latitude),
+            longitude_offset=xarray_utils.infer_longitude_offset(full_era5.longitude),
+        )
+        regridder = horizontal_interpolation.ConservativeRegridder(
+            era5_grid, model.data_coords.horizontal, skipna=True
+        )
+        eval_era5 = xarray_utils.regrid(sliced_era5, regridder)
+        eval_era5 = xarray_utils.fill_nan_with_nearest(eval_era5)
+
+        inner_steps = 24  # save model outputs once every 24 hours
+        outer_steps = 4 * 24 // inner_steps  # total of 4 days
+        timedelta = np.timedelta64(1, 'h') * inner_steps
+        # times = (np.arange(outer_steps) * inner_steps)  # time axis in hours
+
+        # initialize model state
+        inputs = model.inputs_from_xarray(eval_era5.isel(time=0))
+        input_forcings = model.forcings_from_xarray(eval_era5.isel(time=0))
+        rng_key = jax.random.key(42)  # optional for deterministic models
+        initial_state = model.encode(inputs, input_forcings, rng_key)
+
+        # use persistence for forcing variables (SST and sea ice cover)
+        all_forcings = model.forcings_from_xarray(eval_era5.head(time=1))
+
+        # make forecast
+        final_state, predictions = model.unroll(
+            initial_state,
+            all_forcings,
+            steps=outer_steps,
+            timedelta=timedelta,
+            start_with_input=True,
+        )
+        # predictions_ds = model.data_to_xarray(predictions, times=times)
+
+        def sub(initial_state, all_forcings):
+            return model.unroll(initial_state,
+            all_forcings,
+            steps=outer_steps,
+            timedelta=timedelta,
+            start_with_input=True,
+            )
+
+        self.sub = sub
+        self.model = model
+        self.eval_era5 = self.eval_era5
+        self.all_forcings = all_forcings
+        self.outer_steps = outer_steps
+
+    def test(self):
+        for name, pipe, _ in pipelines:
+            if pipe is None:
+                nfn = jax.jit(sub)
+            else:
+                nfn = jax.jit(enzyme_jax_ir(pipeline_options=pipe)(self.sub))
+
+            res = self.run(nfn)
+            print(name, "\t", res)
+
+    def run(self, fn, steps=1):
+        inputs = self.model.inputs_from_xarray(self.eval_era5.isel(time=0))
+        input_forcings = self.model.forcings_from_xarray(self.eval_era5.isel(time=0))
+        rng_key = jax.random.key(42)  # optional for deterministic models
+        initial_state = self.model.encode(inputs, input_forcings, rng_key)
+        map(lambda x:x.block_until_ready(), fn(
+            initial_state,
+            self.all_forcings,
+        ))
+        return timeit.Timer(
+    """map(lambda x:x.block_until_ready(), fn(
         initial_state,
         all_forcings,
-    ))
-    return timeit.Timer(
-"""map(lambda x:x.block_until_ready(), fn(
-    initial_state,
-    all_forcings,
-))""",
-        globals={
-            "fn": fn,
-            "initial_state":initial_state,
-            "timedelta":timedelta,
-            "all_forcings":all_forcings,
-            "outer_steps":outer_steps
-        }
-    ).timeit(steps)
-
-run(jax.jit(sub))
-
-run(jax.jit(enzyme_jax_ir(pipeline_options=JaXPipeline(hlo_opts()))(sub)))
+    ))""",
+            globals={
+                "fn": fn,
+                "initial_state":initial_state,
+                "all_forcings":self.all_forcings,
+            }
+        ).timeit(steps)

From ff7c2679b69a7e3eff4c2f9710d384276125d32a Mon Sep 17 00:00:00 2001
From: "William S. Moses" <gh@wsmoses.com>
Date: Wed, 16 Oct 2024 01:18:56 -0400
Subject: [PATCH 03/17] fixup

---
 test/bench_vs_xla.py |  32 ++++++-------
 test/jaxmd.py        |   2 +-
 test/llama.py        |   3 +-
 test/neuralgcm.py    |  76 +++++++++++++++++-------------
 test/test_utils.py   | 108 +++++++++++++++++++++++++++++++++++--------
 5 files changed, 151 insertions(+), 70 deletions(-)

diff --git a/test/bench_vs_xla.py b/test/bench_vs_xla.py
index b0c1171e..0a472f3c 100644
--- a/test/bench_vs_xla.py
+++ b/test/bench_vs_xla.py
@@ -16,7 +16,7 @@ def setUp(self):
             jnp.array([0.1, 0.2, 0.3]),
             jnp.array([50.0, 70.0, 110.0]),
         ]
-        self.douts = [jnp.array([500.0, 700.0, 110.0])]
+        self.douts = jnp.array([500.0, 700.0, 110.0])
 
         def add_one(x, y):
             return x + 1 + y
@@ -42,7 +42,7 @@ def setUp(self):
             jnp.array([50.0, 70.0, 110.0]),
             jnp.array([1300.0, 1700.0, 1900.0]),
         ]
-        self.douts = [jnp.array([500.0, 700.0, 110.0])]
+        self.douts = jnp.array([500.0, 700.0, 110.0])
 
         def add_two(x, z, y):
             return x + y
@@ -55,7 +55,7 @@ class Sum(EnzymeJaxTest):
     def setUp(self):
         self.ins = [jnp.array(range(50), dtype=jnp.float32)]
         self.dins = [jnp.array([i * i for i in range(50)], dtype=jnp.float32)]
-        self.douts = [1.0]
+        self.douts = jnp.array(1.0)
 
         def sum(x):
             return jnp.sum(x)
@@ -69,7 +69,7 @@ def setUp(self):
         dim = 288
         self.ins = [jnp.array(range(dim), dtype=jnp.float32)]
         self.dins = [jnp.array([i * i for i in range(dim)], dtype=jnp.float32)]
-        self.douts = [jnp.array([i * i for i in range(dim)], dtype=jnp.float32)]
+        self.douts = jnp.array([i * i for i in range(dim)], dtype=jnp.float32)
 
         self.primfilter = no_newxla
         self.fwdfilter = no_newxla
@@ -89,7 +89,7 @@ def setUp(self):
         self.dins = [
             jnp.array([i * i for i in range(dim)], dtype=jnp.float32).reshape(1, dim, 1)
         ]
-        self.douts = [jnp.array([i * i for i in range(dim)], dtype=jnp.float32)]
+        self.douts = jnp.array([i * i for i in range(dim)], dtype=jnp.float32)
 
         self.primfilter = no_newxla
         self.fwdfilter = no_newxla
@@ -107,11 +107,9 @@ def setUp(self):
         dim = 12
         self.ins = [jnp.array(range(dim), dtype=jnp.float32)]
         self.dins = [jnp.array([i * i for i in range(dim)], dtype=jnp.float32)]
-        self.douts = [
-            jnp.array([i * i for i in range(2 * dim)], dtype=jnp.float32).reshape(
-                (2, dim)
-            )
-        ]
+        self.douts = jnp.array(
+            [i * i for i in range(2 * dim)], dtype=jnp.float32
+        ).reshape((2, dim))
 
         self.primfilter = no_newxla
         self.fwdfilter = no_newxla
@@ -134,11 +132,9 @@ def setUp(self):
         dim = 12
         self.ins = [jnp.array(range(dim), dtype=jnp.float32)]
         self.dins = [jnp.array([i * i for i in range(dim)], dtype=jnp.float32)]
-        self.douts = [
-            jnp.array([i * i for i in range(2 * dim)], dtype=jnp.float32).reshape(
-                (2, dim)
-            )
-        ]
+        self.douts = jnp.array(
+            [i * i for i in range(2 * dim)], dtype=jnp.float32
+        ).reshape((2, dim))
 
         self.primfilter = no_newxla
         self.fwdfilter = no_newxla
@@ -174,7 +170,7 @@ def setUp(self):
             jnp.array([i * i for i in range(dim)], dtype=jnp.float32),
             jnp.array([i * i * i / 3.0 for i in range(dim)], dtype=jnp.float32),
         ]
-        self.douts = [jnp.array([i * i for i in range(2 * dim)], dtype=jnp.float32)]
+        self.douts = jnp.array([i * i for i in range(2 * dim)], dtype=jnp.float32)
 
         self.revfilter = justjax
         # self.revfilter = nomlir
@@ -274,7 +270,7 @@ def forward(c_tau):
             jnp.array([2.7, 2.7, 2.7]),
         ]
         self.dins = [jnp.array([3.1, 3.1, 3.1])]
-        self.douts = self.dins
+        self.douts = (self.dins[0],)
         self.revfilter = lambda _: []
         # No support for stablehlo.while atm
         # self.revfilter = justjax
@@ -298,7 +294,7 @@ def forward(position):
 
         self.ins = [jnp.array([2.0, 4.0, 6.0, 8.0])]
         self.dins = [jnp.array([2.7, 3.1, 5.9, 4.2])]
-        self.douts = []
+        self.douts = self.fn(*self.ins)
         self.revfilter = lambda _: []
         # No support for stablehlo.scatter atm
         self.mlirad_rev = False
diff --git a/test/jaxmd.py b/test/jaxmd.py
index 0a4a22bc..40243544 100644
--- a/test/jaxmd.py
+++ b/test/jaxmd.py
@@ -146,7 +146,7 @@ def forward(
         # for i, v in enumerate(self.ins):
         #    print("i=", i, v)
         self.dins = [x.copy() for x in self.ins]
-        self.douts = [x.copy() for x in self.ins]
+        self.douts = tuple(x.copy() for x in self.ins)
         self.AllPipelines = pipelines
         # No support for stablehlo.while atm
         # self.revfilter = justjax
diff --git a/test/llama.py b/test/llama.py
index 84b1fdf4..b0a4e2f1 100644
--- a/test/llama.py
+++ b/test/llama.py
@@ -247,6 +247,7 @@ def forward(x, config, weights, key_cache, value_cache):
 
     return x
 
+
 pipelines = [
     ("JaX  ", None, CurBackends),
     ("JaXPipe", JaXPipeline(), CurBackends),
@@ -336,7 +337,7 @@ def sfn(x, weights, key_cache, value_cache):
 
         self.ins = [x, weights, key_cache, value_cache]
         self.dins = [dx, weights, key_cache, value_cache]
-        self.douts = [dx]
+        self.douts = dx
         self.tol = 5e-5
 
 
diff --git a/test/neuralgcm.py b/test/neuralgcm.py
index e2d58d5e..e8e51181 100644
--- a/test/neuralgcm.py
+++ b/test/neuralgcm.py
@@ -50,39 +50,41 @@
 from dinosaur import xarray_utils
 import neuralgcm
 
+
 class NeuralGCM(absltest.TestCase):
     def setUp(self):
-        gcs = gcsfs.GCSFileSystem(token='anon')
+        gcs = gcsfs.GCSFileSystem(token="anon")
 
-        model_name = 'neural_gcm_dynamic_forcing_deterministic_1_4_deg.pkl'  #@param ['neural_gcm_dynamic_forcing_deterministic_0_7_deg.pkl', 'neural_gcm_dynamic_forcing_deterministic_1_4_deg.pkl', 'neural_gcm_dynamic_forcing_deterministic_2_8_deg.pkl', 'neural_gcm_dynamic_forcing_stochastic_1_4_deg.pkl'] {type: "string"}
+        model_name = "neural_gcm_dynamic_forcing_deterministic_1_4_deg.pkl"  # @param ['neural_gcm_dynamic_forcing_deterministic_0_7_deg.pkl', 'neural_gcm_dynamic_forcing_deterministic_1_4_deg.pkl', 'neural_gcm_dynamic_forcing_deterministic_2_8_deg.pkl', 'neural_gcm_dynamic_forcing_stochastic_1_4_deg.pkl'] {type: "string"}
 
-        with gcs.open(f'gs://gresearch/neuralgcm/04_30_2024/{model_name}', 'rb') as f:
-          ckpt = pickle.load(f)
+        with gcs.open(f"gs://gresearch/neuralgcm/04_30_2024/{model_name}", "rb") as f:
+            ckpt = pickle.load(f)
 
         model = neuralgcm.PressureLevelModel.from_checkpoint(ckpt)
 
-        era5_path = 'gs://gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3'
+        era5_path = (
+            "gs://gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3"
+        )
         full_era5 = xarray.open_zarr(gcs.get_mapper(era5_path), chunks=None)
 
-        demo_start_time = '2020-02-14'
-        demo_end_time = '2020-02-18'
+        demo_start_time = "2020-02-14"
+        demo_end_time = "2020-02-18"
         data_inner_steps = 24  # process every 24th hour
 
         sliced_era5 = (
-            full_era5
-            [model.input_variables + model.forcing_variables]
+            full_era5[model.input_variables + model.forcing_variables]
             .pipe(
                 xarray_utils.selective_temporal_shift,
                 variables=model.forcing_variables,
-                time_shift='24 hours',
+                time_shift="24 hours",
             )
             .sel(time=slice(demo_start_time, demo_end_time, data_inner_steps))
             .compute()
         )
 
         era5_grid = spherical_harmonic.Grid(
-            latitude_nodes=full_era5.sizes['latitude'],
-            longitude_nodes=full_era5.sizes['longitude'],
+            latitude_nodes=full_era5.sizes["latitude"],
+            longitude_nodes=full_era5.sizes["longitude"],
             latitude_spacing=xarray_utils.infer_latitude_spacing(full_era5.latitude),
             longitude_offset=xarray_utils.infer_longitude_offset(full_era5.longitude),
         )
@@ -92,9 +94,12 @@ def setUp(self):
         eval_era5 = xarray_utils.regrid(sliced_era5, regridder)
         eval_era5 = xarray_utils.fill_nan_with_nearest(eval_era5)
 
-        inner_steps = 24  # save model outputs once every 24 hours
-        outer_steps = 4 * 24 // inner_steps  # total of 4 days
-        timedelta = np.timedelta64(1, 'h') * inner_steps
+        # inner_steps = 24  # save model outputs once every 24 hours
+        # outer_steps = 4 * 24 // inner_steps  # total of 4 days
+        inner_steps = 4  # save model outputs once every 24 hours
+        outer_steps = 4 * 4 // inner_steps  # total of 4 days
+
+        timedelta = np.timedelta64(1, "h") * inner_steps
         # times = (np.arange(outer_steps) * inner_steps)  # time axis in hours
 
         # initialize model state
@@ -117,11 +122,12 @@ def setUp(self):
         # predictions_ds = model.data_to_xarray(predictions, times=times)
 
         def sub(initial_state, all_forcings):
-            return model.unroll(initial_state,
-            all_forcings,
-            steps=outer_steps,
-            timedelta=timedelta,
-            start_with_input=True,
+            return model.unroll(
+                initial_state,
+                all_forcings,
+                steps=outer_steps,
+                timedelta=timedelta,
+                start_with_input=True,
             )
 
         self.sub = sub
@@ -132,31 +138,39 @@ def sub(initial_state, all_forcings):
 
     def test(self):
         for name, pipe, _ in pipelines:
+            print("name=", name, res)
             if pipe is None:
                 nfn = jax.jit(sub)
             else:
                 nfn = jax.jit(enzyme_jax_ir(pipeline_options=pipe)(self.sub))
 
-            res = self.run(nfn)
-            print(name, "\t", res)
+            res = self.run_on_fn(nfn)
+            print("name=", name, res)
 
-    def run(self, fn, steps=1):
+    def run_on_fn(self, fn, steps=1):
         inputs = self.model.inputs_from_xarray(self.eval_era5.isel(time=0))
         input_forcings = self.model.forcings_from_xarray(self.eval_era5.isel(time=0))
         rng_key = jax.random.key(42)  # optional for deterministic models
         initial_state = self.model.encode(inputs, input_forcings, rng_key)
-        map(lambda x:x.block_until_ready(), fn(
-            initial_state,
-            self.all_forcings,
-        ))
+        map(
+            lambda x: x.block_until_ready(),
+            fn(
+                initial_state,
+                self.all_forcings,
+            ),
+        )
         return timeit.Timer(
-    """map(lambda x:x.block_until_ready(), fn(
+            """map(lambda x:x.block_until_ready(), fn(
         initial_state,
         all_forcings,
     ))""",
             globals={
                 "fn": fn,
-                "initial_state":initial_state,
-                "all_forcings":self.all_forcings,
-            }
+                "initial_state": initial_state,
+                "all_forcings": self.all_forcings,
+            },
         ).timeit(steps)
+
+
+if __name__ == "__main__":
+    absltest.main()
diff --git a/test/test_utils.py b/test/test_utils.py
index c95f6f96..0054e37b 100644
--- a/test/test_utils.py
+++ b/test/test_utils.py
@@ -145,7 +145,6 @@ def setup_backends():
 )
 
 
-
 def no_newxla(x):
     return [
         (name, a, b) for (name, a, b) in x if name != "NewXLAMLIR" and name != "NewXLA"
@@ -183,6 +182,26 @@ def fwd(*args):
     return fwd
 
 
+def sync(x):
+    return x.block_until_ready()
+
+
+def syncall(x):
+    return map(sync, x)
+
+
+def fwdsync1(x):
+    return (sync(x[0]), sync(x[1]))
+
+
+def fwdsync2(x):
+    return (sync(x[0][0]), sync(x[1][0]))
+
+
+def fwdsync3(x):
+    return (syncall(x[0]), syncall(x[1]))
+
+
 # @jax.jit
 # def rev_jax(dout, in0, in1):
 # primals, f_vjp = jax.vjp(add_one_jax, in0, in1)
@@ -201,11 +220,27 @@ def splatvjp_noprim(in_fn):
     def rev(dout, *args):
         primals, f_vjp = jax.vjp(in_fn, *args)
         grads = f_vjp(dout)
-        return primals, grads
+        return grads
 
     return rev
 
 
+def revsync0_0(x):
+    return (sync(x[0]), sync(x[1][0]))
+
+
+def revsync0_1(x):
+    return (syncall(x[0]), sync(x[1][0]))
+
+
+def revsync1_0(x):
+    return (sync(x[0]), syncall(x[1]))
+
+
+def revsync1_1(x):
+    return (syncall(x[0]), syncall(x[1]))
+
+
 def to_backend(x, backend):
     dev = jax.local_devices(backend=backend)[0]
     return jax.device_put(x, dev)
@@ -266,6 +301,12 @@ def harness(self, name, in_fn, ins, dins, douts):
         assert len(ins) == len(dins)
 
         primalstr = "fn(" + (", ".join(["in" + str(i) for i in range(len(ins))])) + ")"
+        if isinstance(douts, jax.Array):
+            primalstr = "sync(" + primalstr + ")"
+        elif len(douts) == 1:
+            primalstr = "sync(" + primalstr + "[0])"
+        else:
+            primalstr = "syncall(" + primalstr + "[0])"
 
         fwdstr = (
             "fwd("
@@ -274,24 +315,57 @@ def harness(self, name, in_fn, ins, dins, douts):
             + (", ".join(["din" + str(i) for i in range(len(dins))]))
             + ")"
         )
-
-        revstr = (
+        if isinstance(douts, jax.Array):
+            fwdstr = "fwdsync1(" + fwdstr + ")"
+        elif len(douts) == 1:
+            fwdstr = "fwdsync2(" + fwdstr + ")"
+        else:
+            fwdstr = "fwdsync3(" + fwdstr + ")"
+
+        revstr0 = (
             "rev(dout, " + (", ".join(["in" + str(i) for i in range(len(ins))])) + ")"
         )
+        if self.revprimal:
+            if len(dins) == 1:
+                if isinstance(douts, jax.Array):
+                    revstr = "revsync0_0(" + revstr0 + ")"
+                else:
+                    revstr = "revsync0_1(" + revstr0 + ")"
+            else:
+                if isinstance(douts, jax.Array):
+                    revstr = "revsync1_0(" + revstr0 + ")"
+                else:
+                    revstr = "revsync1_1(" + revstr0 + ")"
+        else:
+            if len(dins) == 1:
+                revstr = "sync(" + revstr0 + "[0])"
+            else:
+                revstr = "syncall(" + revstr0 + ")"
+        revtransform = splatvjp if self.revprimal else splatvjp_noprim
 
         for backend in self.AllBackends:
             ins_backend = [to_backend(x, backend) for x in ins]
             dins_backend = [to_backend(x, backend) for x in dins]
-            douts_backend = [to_backend(x, backend) for x in douts]
+            douts_backend = None
+            if isinstance(douts, jax.Array):
+                douts_backend = to_backend(douts, backend)
+            else:
+                douts_backend = tuple([to_backend(x, backend) for x in douts])
 
             primalins = {("in" + str(i)): ins_backend[i] for i in range(len(ins))}
+            primalins["sync"] = sync
+            primalins["syncall"] = syncall
+            primalins["fwdsync1"] = fwdsync1
+            primalins["fwdsync2"] = fwdsync2
+            primalins["fwdsync3"] = fwdsync3
+            primalins["revsync0_0"] = revsync0_0
+            primalins["revsync0_1"] = revsync0_1
+            primalins["revsync1_0"] = revsync1_0
+            primalins["revsync1_1"] = revsync1_1
             fwdins = primalins | {
                 ("din" + str(i)): dins_backend[i] for i in range(len(dins))
             }
-            if len(douts) == 1:
-                revins = primalins | {"dout": douts_backend[0]}
-            else:
-                revins = primalins | {"dout": tuple(douts_backend)}
+            revins = primalins | {"dout": douts_backend}
             primres = None
 
             for pname, pipeline, pbackends in self.primfilter(self.AllPipelines):
@@ -386,14 +460,10 @@ def harness(self, name, in_fn, ins, dins, douts):
 
             revres = None
 
-            revtransform = splatvjp if self.revprimal else splatvjp_noprim
-
             for pname, pipeline, pbackends in self.revfilter(self.AllPipelines):
                 if backend in pbackends:
 
                     adout = douts_backend
-                    if len(douts) != 1:
-                        adout = [tuple(douts_backend)]
                     if pipeline is not None:
                         if self.mlirad_rev or pipeline is None:
                             rfn_enzyme = (
@@ -409,9 +479,9 @@ def harness(self, name, in_fn, ins, dins, douts):
                             )
 
                             if self.revprimal:
-                                primals, grads = rev_enzyme(*adout, *ins_backend)
+                                primals, grads = rev_enzyme(adout, *ins_backend)
                             else:
-                                grads = rev_enzyme(*adout, *ins_backend)
+                                grads = rev_enzyme(adout, *ins_backend)
                                 assert grads is not None
 
                             if self.revprimal and primres is not None:
@@ -455,9 +525,9 @@ def harness(self, name, in_fn, ins, dins, douts):
                         )
 
                         if self.revprimal:
-                            primals, grads = rev_enzyme(*adout, *ins_backend)
+                            primals, grads = rev_enzyme(adout, *ins_backend)
                         else:
-                            grads = rev_enzyme(*adout, *ins_backend)
+                            grads = rev_enzyme(adout, *ins_backend)
                             assert grads is not None
 
                         if self.revprimal and primres is not None:
@@ -508,9 +578,9 @@ def harness(self, name, in_fn, ins, dins, douts):
                         )
 
                         if self.revprimal:
-                            primals, grads = rev_enzyme(*adout, *ins_backend)
+                            primals, grads = rev_enzyme(adout, *ins_backend)
                         else:
-                            grads = rev_enzyme(*adout, *ins_backend)
+                            grads = rev_enzyme(adout, *ins_backend)
                             assert grads is not None
 
                         if self.revprimal and primres is not None:

From 3d99c792a141bd384d5c246370362cc0e4e3e73f Mon Sep 17 00:00:00 2001
From: "William S. Moses" <gh@wsmoses.com>
Date: Wed, 16 Oct 2024 01:32:03 -0400
Subject: [PATCH 04/17] fixup

---
 test/BUILD        | 2 +-
 test/neuralgcm.py | 8 ++++----
 2 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/test/BUILD b/test/BUILD
index 9fda413f..6f514a24 100644
--- a/test/BUILD
+++ b/test/BUILD
@@ -144,7 +144,7 @@ py_test(
     ],
     imports = ["."],
     deps = TEST_DEPS + ["@pypi_neuralgcm//:pkg", "@pypi_gcsfs//:pkg"],
-    timeout='long'
+    timeout='eternal'
 )
 
 
diff --git a/test/neuralgcm.py b/test/neuralgcm.py
index e8e51181..e98d2652 100644
--- a/test/neuralgcm.py
+++ b/test/neuralgcm.py
@@ -94,10 +94,10 @@ def setUp(self):
         eval_era5 = xarray_utils.regrid(sliced_era5, regridder)
         eval_era5 = xarray_utils.fill_nan_with_nearest(eval_era5)
 
-        # inner_steps = 24  # save model outputs once every 24 hours
-        # outer_steps = 4 * 24 // inner_steps  # total of 4 days
-        inner_steps = 4  # save model outputs once every 24 hours
-        outer_steps = 4 * 4 // inner_steps  # total of 4 days
+        inner_steps = 24  # save model outputs once every 24 hours
+        outer_steps = 4 * 24 // inner_steps  # total of 4 days
+        # inner_steps = 4  # save model outputs once every 24 hours
+        # outer_steps = 4 * 4 // inner_steps  # total of 4 days
 
         timedelta = np.timedelta64(1, "h") * inner_steps
         # times = (np.arange(outer_steps) * inner_steps)  # time axis in hours

From 0eb6eba753d0123710fa6f1bf6ea329063fb75b7 Mon Sep 17 00:00:00 2001
From: "William S. Moses" <gh@wsmoses.com>
Date: Wed, 16 Oct 2024 01:34:18 -0400
Subject: [PATCH 05/17] fixup

---
 test/neuralgcm.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/test/neuralgcm.py b/test/neuralgcm.py
index e98d2652..ada1b6c5 100644
--- a/test/neuralgcm.py
+++ b/test/neuralgcm.py
@@ -132,7 +132,7 @@ def sub(initial_state, all_forcings):
 
         self.sub = sub
         self.model = model
-        self.eval_era5 = self.eval_era5
+        self.eval_era5 = eval_era5
         self.all_forcings = all_forcings
         self.outer_steps = outer_steps
 

From aa12ed04a9cf985710f7f68e9a1261b11639d34a Mon Sep 17 00:00:00 2001
From: "William S. Moses" <gh@wsmoses.com>
Date: Wed, 16 Oct 2024 01:56:00 -0400
Subject: [PATCH 06/17] fix

---
 test/neuralgcm.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/test/neuralgcm.py b/test/neuralgcm.py
index ada1b6c5..4356009c 100644
--- a/test/neuralgcm.py
+++ b/test/neuralgcm.py
@@ -138,9 +138,9 @@ def sub(initial_state, all_forcings):
 
     def test(self):
         for name, pipe, _ in pipelines:
-            print("name=", name, res)
+            print("name=", name)
             if pipe is None:
-                nfn = jax.jit(sub)
+                nfn = jax.jit(self.sub)
             else:
                 nfn = jax.jit(enzyme_jax_ir(pipeline_options=pipe)(self.sub))
 

From d84c69a261acb544ae12ea2ae855e0823207c520 Mon Sep 17 00:00:00 2001
From: "William S. Moses" <gh@wsmoses.com>
Date: Wed, 16 Oct 2024 10:52:58 -0400
Subject: [PATCH 07/17] create state once

---
 test/neuralgcm.py | 13 +++++++------
 1 file changed, 7 insertions(+), 6 deletions(-)

diff --git a/test/neuralgcm.py b/test/neuralgcm.py
index 4356009c..e0177b82 100644
--- a/test/neuralgcm.py
+++ b/test/neuralgcm.py
@@ -135,6 +135,11 @@ def sub(initial_state, all_forcings):
         self.eval_era5 = eval_era5
         self.all_forcings = all_forcings
         self.outer_steps = outer_steps
+        
+        inputs = self.model.inputs_from_xarray(self.eval_era5.isel(time=0))
+        input_forcings = self.model.forcings_from_xarray(self.eval_era5.isel(time=0))
+        rng_key = jax.random.key(42)  # optional for deterministic models
+        self.initial_state = self.model.encode(inputs, input_forcings, rng_key)
 
     def test(self):
         for name, pipe, _ in pipelines:
@@ -148,14 +153,10 @@ def test(self):
             print("name=", name, res)
 
     def run_on_fn(self, fn, steps=1):
-        inputs = self.model.inputs_from_xarray(self.eval_era5.isel(time=0))
-        input_forcings = self.model.forcings_from_xarray(self.eval_era5.isel(time=0))
-        rng_key = jax.random.key(42)  # optional for deterministic models
-        initial_state = self.model.encode(inputs, input_forcings, rng_key)
         map(
             lambda x: x.block_until_ready(),
             fn(
-                initial_state,
+                self.initial_state,
                 self.all_forcings,
             ),
         )
@@ -166,7 +167,7 @@ def run_on_fn(self, fn, steps=1):
     ))""",
             globals={
                 "fn": fn,
-                "initial_state": initial_state,
+                "initial_state": self.initial_state,
                 "all_forcings": self.all_forcings,
             },
         ).timeit(steps)

From d68a6beccb9eabe9b74e13668fee2557d65b2571 Mon Sep 17 00:00:00 2001
From: "William S. Moses" <gh@wsmoses.com>
Date: Wed, 16 Oct 2024 14:15:09 -0400
Subject: [PATCH 08/17] cleanup

---
 test/neuralgcm.py | 15 ++++++++-------
 1 file changed, 8 insertions(+), 7 deletions(-)

diff --git a/test/neuralgcm.py b/test/neuralgcm.py
index e0177b82..6babf14d 100644
--- a/test/neuralgcm.py
+++ b/test/neuralgcm.py
@@ -1,4 +1,3 @@
-from absl.testing import absltest
 import jax.numpy as jnp
 import jax.random
 import jax.lax
@@ -51,7 +50,7 @@
 import neuralgcm
 
 
-class NeuralGCM(absltest.TestCase):
+class NeuralGCM:
     def setUp(self):
         gcs = gcsfs.GCSFileSystem(token="anon")
 
@@ -94,10 +93,10 @@ def setUp(self):
         eval_era5 = xarray_utils.regrid(sliced_era5, regridder)
         eval_era5 = xarray_utils.fill_nan_with_nearest(eval_era5)
 
-        inner_steps = 24  # save model outputs once every 24 hours
-        outer_steps = 4 * 24 // inner_steps  # total of 4 days
-        # inner_steps = 4  # save model outputs once every 24 hours
-        # outer_steps = 4 * 4 // inner_steps  # total of 4 days
+        # inner_steps = 24  # save model outputs once every 24 hours
+        # outer_steps = 4 * 24 // inner_steps  # total of 4 days
+        inner_steps = 4  # save model outputs once every 24 hours
+        outer_steps = 4 * 4 // inner_steps  # total of 4 days
 
         timedelta = np.timedelta64(1, "h") * inner_steps
         # times = (np.arange(outer_steps) * inner_steps)  # time axis in hours
@@ -174,4 +173,6 @@ def run_on_fn(self, fn, steps=1):
 
 
 if __name__ == "__main__":
-    absltest.main()
+    c = NeuralGCM()
+    c.setUp()
+    c.test()

From 0c19703a6e633a14b150803a1c449d2cecf11222 Mon Sep 17 00:00:00 2001
From: William Moses <gh@wsmoses.com>
Date: Fri, 18 Oct 2024 15:24:23 -0700
Subject: [PATCH 09/17] Update maxtext.py

---
 test/maxtext.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/test/maxtext.py b/test/maxtext.py
index 2f2a9267..de700e81 100644
--- a/test/maxtext.py
+++ b/test/maxtext.py
@@ -1,9 +1,8 @@
 # Steps for getting results here
 # Run:
-#   1) pip install https://github.com/wsmoses/maxtext
-#   2) bazel build -c opt //:wheel
-#   3) pip install ./bazel-bin/*whl
-#   4) python test/maxtext.py
+#   1) bazel build -c opt //:wheel
+#   2) python3 -m pip install --user ./bazel-bin/*.whl  git+https://github.com/wsmoses/maxtext aqtp tensorboardX google-cloud-storage datasets neuralgcm gcsfs --break-system-packages
+#   3) python test/maxtext.py
 
 from absl.testing import absltest
 import jax.numpy as jnp
@@ -19,6 +18,7 @@
 )
 import numpy as np
 import timeit
+from test_utils import partialopt
 
 argv = ("-I/usr/include/c++/11", "-I/usr/include/x86_64-linux-gnu/c++/11")
 

From 38fed921e34cf0522e28fa9453deb048e0c794af Mon Sep 17 00:00:00 2001
From: "William S. Moses" <wmoses@google.com>
Date: Wed, 6 Nov 2024 21:52:38 +0000
Subject: [PATCH 10/17] fix

---
 test/{neuralgcm.py => neuralgcm_test.py} | 0
 1 file changed, 0 insertions(+), 0 deletions(-)
 rename test/{neuralgcm.py => neuralgcm_test.py} (100%)

diff --git a/test/neuralgcm.py b/test/neuralgcm_test.py
similarity index 100%
rename from test/neuralgcm.py
rename to test/neuralgcm_test.py

From b5c9e5d9c1196ae2f30ab95288c1875706cdda22 Mon Sep 17 00:00:00 2001
From: "William S. Moses" <gh@wsmoses.com>
Date: Thu, 7 Nov 2024 14:24:45 -0500
Subject: [PATCH 11/17] fix ci

---
 .buildkite/gpu_pipeline.yml | 13 ++++---------
 1 file changed, 4 insertions(+), 9 deletions(-)

diff --git a/.buildkite/gpu_pipeline.yml b/.buildkite/gpu_pipeline.yml
index 2a52f2b8..e31a3712 100644
--- a/.buildkite/gpu_pipeline.yml
+++ b/.buildkite/gpu_pipeline.yml
@@ -63,23 +63,18 @@ steps:
       HERMETIC_PYTHON_VERSION="3.12" .local/bin/bazel --output_user_root=`pwd`/.baztmp test --repo_env CUDA_DIR --repo_env XLA_FLAGS --action_env XLA_FLAGS --repo_env TF_CPP_MIN_LOG_LEVEL --action_env TF_CPP_MIN_LOG_LEVEL --test_output=errors //test/... || echo "fail1"
       HERMETIC_PYTHON_VERSION="3.12" .local/bin/bazel --output_user_root=`pwd`/.baztmp test --repo_env CUDA_DIR --repo_env XLA_FLAGS --action_env XLA_FLAGS --repo_env TF_CPP_MIN_LOG_LEVEL --action_env TF_CPP_MIN_LOG_LEVEL --cache_test_results=no -s //test:bench_vs_xla 
       HERMETIC_PYTHON_VERSION="3.12" .local/bin/bazel --output_user_root=`pwd`/.baztmp test --repo_env CUDA_DIR --repo_env XLA_FLAGS --action_env XLA_FLAGS --repo_env TF_CPP_MIN_LOG_LEVEL --action_env TF_CPP_MIN_LOG_LEVEL --cache_test_results=no -s //test:llama 
-      HERMETIC_PYTHON_VERSION="3.12" .local/bin/bazel --output_user_root=`pwd`/.baztmp test --repo_env CUDA_DIR --repo_env XLA_FLAGS --action_env XLA_FLAGS --repo_env TF_CPP_MIN_LOG_LEVEL --action_env TF_CPP_MIN_LOG_LEVEL --cache_test_results=no -s //test:jaxmd
-      
-      if [ -f "test/test_neuralgcm.py" ]; then
-        HERMETIC_PYTHON_VERSION="3.12" .local/bin/bazel --output_user_root=`pwd`/.baztmp test --repo_env CUDA_DIR --repo_env XLA_FLAGS --action_env XLA_FLAGS --repo_env TF_CPP_MIN_LOG_LEVEL --action_env TF_CPP_MIN_LOG_LEVEL --cache_test_results=no -s //test:test_neuralgcm
-      fi
+      HERMETIC_PYTHON_VERSION="3.12" .local/bin/bazel --output_user_root=`pwd`/.baztmp test --repo_env CUDA_DIR --repo_env XLA_FLAGS --action_env XLA_FLAGS --repo_env TF_CPP_MIN_LOG_LEVEL --action_env TF_CPP_MIN_LOG_LEVEL --cache_test_results=no -s //test:jaxmd 
+      HERMETIC_PYTHON_VERSION="3.12" .local/bin/bazel --output_user_root=`pwd`/.baztmp test --repo_env CUDA_DIR --repo_env XLA_FLAGS --action_env XLA_FLAGS --repo_env TF_CPP_MIN_LOG_LEVEL --action_env TF_CPP_MIN_LOG_LEVEL --cache_test_results=no -s //test:neuralgcm_test
       bazel-bin/test/llama.runfiles/python_*/bin/python3 -m pip install bazel-bin/*.whl https://github.com/wsmoses/maxtext aqtp tensorboardX google-cloud-storage datasets gcsfs
       bazel-bin/test/llama.runfiles/python_*/bin/python3 test/maxtext.py > maxtext.log
       cat bazel-out/*/testlogs/test/llama/test.log
       cat bazel-out/*/testlogs/test/bench_vs_xla/test.log
       cat bazel-out/*/testlogs/test/jaxmd/test.log
-      if [ -f "test/test_neuralgcm.py" ]; then
-        cat bazel-out/*/testlogs/test/test_neuralgcm/test.log
-      fi
+      cat bazel-out/*/testlogs/test/neuralgcm_test/test.log
       cat maxtext.log
     artifact_paths:
       - "bazel-out/*/testlogs/test/llama/test.log"
       - "bazel-out/*/testlogs/test/bench_vs_xla/test.log"
       - "bazel-out/*/testlogs/test/jaxmd/test.log"
-      - "bazel-out/*/testlogs/test/test_neuralgcm/test.log"
+      - "bazel-out/*/testlogs/test/neuralgcm_test/test.log"
       - "maxtext.log"

From dd1301722c00a765f128927dd8d20e0508abcb0e Mon Sep 17 00:00:00 2001
From: "William S. Moses" <gh@wsmoses.com>
Date: Thu, 7 Nov 2024 14:33:04 -0500
Subject: [PATCH 12/17] fmt

---
 test/neuralgcm_test.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/test/neuralgcm_test.py b/test/neuralgcm_test.py
index 6babf14d..9db4aac1 100644
--- a/test/neuralgcm_test.py
+++ b/test/neuralgcm_test.py
@@ -134,7 +134,7 @@ def sub(initial_state, all_forcings):
         self.eval_era5 = eval_era5
         self.all_forcings = all_forcings
         self.outer_steps = outer_steps
-        
+
         inputs = self.model.inputs_from_xarray(self.eval_era5.isel(time=0))
         input_forcings = self.model.forcings_from_xarray(self.eval_era5.isel(time=0))
         rng_key = jax.random.key(42)  # optional for deterministic models

From 8628241fd123a5d5e1200444630349e045d98343 Mon Sep 17 00:00:00 2001
From: "William S. Moses" <gh@wsmoses.com>
Date: Thu, 7 Nov 2024 16:44:19 -0500
Subject: [PATCH 13/17] fix

---
 builddeps/requirements_lock_3_10.txt | 879 ++++++++++++++++++++++++++-
 builddeps/requirements_lock_3_11.txt | 871 +++++++++++++++++++++++++-
 builddeps/requirements_lock_3_12.txt | 874 +++++++++++++++++++++++++-
 test/BUILD                           |   4 +-
 4 files changed, 2601 insertions(+), 27 deletions(-)

diff --git a/builddeps/requirements_lock_3_10.txt b/builddeps/requirements_lock_3_10.txt
index 7d223b9b..3dd86c8a 100644
--- a/builddeps/requirements_lock_3_10.txt
+++ b/builddeps/requirements_lock_3_10.txt
@@ -16,10 +16,125 @@ absl-py==2.1.0 \
     #   ml-collections
     #   optax
     #   orbax-checkpoint
+    #   tensorflow-probability
+aiohappyeyeballs==2.4.3 \
+    --hash=sha256:75cf88a15106a5002a8eb1dab212525c00d1f4c0fa96e551c9fbe6f09a621586 \
+    --hash=sha256:8a7a83727b2756f394ab2895ea0765a0a8c475e3c71e98d43d76f22b4b435572
+    # via aiohttp
+aiohttp==3.10.10 \
+    --hash=sha256:007ec22fbc573e5eb2fb7dec4198ef8f6bf2fe4ce20020798b2eb5d0abda6138 \
+    --hash=sha256:00819de9e45d42584bed046314c40ea7e9aea95411b38971082cad449392b08c \
+    --hash=sha256:01948b1d570f83ee7bbf5a60ea2375a89dfb09fd419170e7f5af029510033d24 \
+    --hash=sha256:038f514fe39e235e9fef6717fbf944057bfa24f9b3db9ee551a7ecf584b5b480 \
+    --hash=sha256:03a42ac7895406220124c88911ebee31ba8b2d24c98507f4a8bf826b2937c7f2 \
+    --hash=sha256:05646ebe6b94cc93407b3bf34b9eb26c20722384d068eb7339de802154d61bc5 \
+    --hash=sha256:0631dd7c9f0822cc61c88586ca76d5b5ada26538097d0f1df510b082bad3411a \
+    --hash=sha256:0b00807e2605f16e1e198f33a53ce3c4523114059b0c09c337209ae55e3823a8 \
+    --hash=sha256:0e1b370d8007c4ae31ee6db7f9a2fe801a42b146cec80a86766e7ad5c4a259cf \
+    --hash=sha256:15ecd889a709b0080f02721255b3f80bb261c2293d3c748151274dfea93ac871 \
+    --hash=sha256:1b66ccafef7336a1e1f0e389901f60c1d920102315a56df85e49552308fc0486 \
+    --hash=sha256:1bbb122c557a16fafc10354b9d99ebf2f2808a660d78202f10ba9d50786384b9 \
+    --hash=sha256:1eb89d3d29adaf533588f209768a9c02e44e4baf832b08118749c5fad191781d \
+    --hash=sha256:258c5dd01afc10015866114e210fb7365f0d02d9d059c3c3415382ab633fcbcb \
+    --hash=sha256:2609e9ab08474702cc67b7702dbb8a80e392c54613ebe80db7e8dbdb79837c68 \
+    --hash=sha256:274cfa632350225ce3fdeb318c23b4a10ec25c0e2c880eff951a3842cf358ac1 \
+    --hash=sha256:28529e08fde6f12eba8677f5a8608500ed33c086f974de68cc65ab218713a59d \
+    --hash=sha256:2b606353da03edcc71130b52388d25f9a30a126e04caef1fd637e31683033abd \
+    --hash=sha256:30ca7c3b94708a9d7ae76ff281b2f47d8eaf2579cd05971b5dc681db8caac6e1 \
+    --hash=sha256:333cf6cf8e65f6a1e06e9eb3e643a0c515bb850d470902274239fea02033e9a8 \
+    --hash=sha256:3455522392fb15ff549d92fbf4b73b559d5e43dc522588f7eb3e54c3f38beee7 \
+    --hash=sha256:362f641f9071e5f3ee6f8e7d37d5ed0d95aae656adf4ef578313ee585b585959 \
+    --hash=sha256:3bcd391d083f636c06a68715e69467963d1f9600f85ef556ea82e9ef25f043f7 \
+    --hash=sha256:3dffb610a30d643983aeb185ce134f97f290f8935f0abccdd32c77bed9388b42 \
+    --hash=sha256:3fe407bf93533a6fa82dece0e74dbcaaf5d684e5a51862887f9eaebe6372cd79 \
+    --hash=sha256:413251f6fcf552a33c981c4709a6bba37b12710982fec8e558ae944bfb2abd38 \
+    --hash=sha256:438cd072f75bb6612f2aca29f8bd7cdf6e35e8f160bc312e49fbecab77c99e3a \
+    --hash=sha256:4470c73c12cd9109db8277287d11f9dd98f77fc54155fc71a7738a83ffcc8ea8 \
+    --hash=sha256:45c3b868724137f713a38376fef8120c166d1eadd50da1855c112fe97954aed8 \
+    --hash=sha256:486f7aabfa292719a2753c016cc3a8f8172965cabb3ea2e7f7436c7f5a22a151 \
+    --hash=sha256:4f05e9727ce409358baa615dbeb9b969db94324a79b5a5cea45d39bdb01d82e6 \
+    --hash=sha256:50aed5155f819873d23520919e16703fc8925e509abbb1a1491b0087d1cd969e \
+    --hash=sha256:50edbcad60d8f0e3eccc68da67f37268b5144ecc34d59f27a02f9611c1d4eec7 \
+    --hash=sha256:54ca74df1be3c7ca1cf7f4c971c79c2daf48d9aa65dea1a662ae18926f5bc8ce \
+    --hash=sha256:578a4b875af3e0daaf1ac6fa983d93e0bbfec3ead753b6d6f33d467100cdc67b \
+    --hash=sha256:597a079284b7ee65ee102bc3a6ea226a37d2b96d0418cc9047490f231dc09fe8 \
+    --hash=sha256:59bb3c54aa420521dc4ce3cc2c3fe2ad82adf7b09403fa1f48ae45c0cbde6628 \
+    --hash=sha256:5c6a5b8c7926ba5d8545c7dd22961a107526562da31a7a32fa2456baf040939f \
+    --hash=sha256:64f6c17757251e2b8d885d728b6433d9d970573586a78b78ba8929b0f41d045a \
+    --hash=sha256:679abe5d3858b33c2cf74faec299fda60ea9de62916e8b67e625d65bf069a3b7 \
+    --hash=sha256:741a46d58677d8c733175d7e5aa618d277cd9d880301a380fd296975a9cdd7bc \
+    --hash=sha256:7789050d9e5d0c309c706953e5e8876e38662d57d45f936902e176d19f1c58ab \
+    --hash=sha256:77abf6665ae54000b98b3c742bc6ea1d1fb31c394bcabf8b5d2c1ac3ebfe7f3b \
+    --hash=sha256:79019094f87c9fb44f8d769e41dbb664d6e8fcfd62f665ccce36762deaa0e911 \
+    --hash=sha256:7b06b7843929e41a94ea09eb1ce3927865387e3e23ebe108e0d0d09b08d25be9 \
+    --hash=sha256:7e338c0523d024fad378b376a79faff37fafb3c001872a618cde1d322400a572 \
+    --hash=sha256:7ea7ffc6d6d6f8a11e6f40091a1040995cdff02cfc9ba4c2f30a516cb2633554 \
+    --hash=sha256:8105fd8a890df77b76dd3054cddf01a879fc13e8af576805d667e0fa0224c35d \
+    --hash=sha256:84afcdea18eda514c25bc68b9af2a2b1adea7c08899175a51fe7c4fb6d551257 \
+    --hash=sha256:9294bbb581f92770e6ed5c19559e1e99255e4ca604a22c5c6397b2f9dd3ee42c \
+    --hash=sha256:93429602396f3383a797a2a70e5f1de5df8e35535d7806c9f91df06f297e109b \
+    --hash=sha256:9627cc1a10c8c409b5822a92d57a77f383b554463d1884008e051c32ab1b3742 \
+    --hash=sha256:998f3bd3cfc95e9424a6acd7840cbdd39e45bc09ef87533c006f94ac47296090 \
+    --hash=sha256:9c72109213eb9d3874f7ac8c0c5fa90e072d678e117d9061c06e30c85b4cf0e6 \
+    --hash=sha256:9fc1500fd2a952c5c8e3b29aaf7e3cc6e27e9cfc0a8819b3bce48cc1b849e4cc \
+    --hash=sha256:a3f00003de6eba42d6e94fabb4125600d6e484846dbf90ea8e48a800430cc142 \
+    --hash=sha256:a45d85cf20b5e0d0aa5a8dca27cce8eddef3292bc29d72dcad1641f4ed50aa16 \
+    --hash=sha256:a7d8d14fe962153fc681f6366bdec33d4356f98a3e3567782aac1b6e0e40109a \
+    --hash=sha256:a8fa23fe62c436ccf23ff930149c047f060c7126eae3ccea005f0483f27b2e28 \
+    --hash=sha256:aa6658732517ddabe22c9036479eabce6036655ba87a0224c612e1ae6af2087e \
+    --hash=sha256:aafc8ee9b742ce75044ae9a4d3e60e3d918d15a4c2e08a6c3c3e38fa59b92d94 \
+    --hash=sha256:ab5a5a0c7a7991d90446a198689c0535be89bbd6b410a1f9a66688f0880ec026 \
+    --hash=sha256:acd48d5b80ee80f9432a165c0ac8cbf9253eaddb6113269a5e18699b33958dbb \
+    --hash=sha256:ad7593bb24b2ab09e65e8a1d385606f0f47c65b5a2ae6c551db67d6653e78c28 \
+    --hash=sha256:baa42524a82f75303f714108fea528ccacf0386af429b69fff141ffef1c534f9 \
+    --hash=sha256:bdfcf6443637c148c4e1a20c48c566aa694fa5e288d34b20fcdc58507882fed3 \
+    --hash=sha256:be7443669ae9c016b71f402e43208e13ddf00912f47f623ee5994e12fc7d4b3f \
+    --hash=sha256:c02a30b904282777d872266b87b20ed8cc0d1501855e27f831320f471d54d983 \
+    --hash=sha256:c1277cd707c465cd09572a774559a3cc7c7a28802eb3a2a9472588f062097205 \
+    --hash=sha256:c30a0eafc89d28e7f959281b58198a9fa5e99405f716c0289b7892ca345fe45f \
+    --hash=sha256:c5ce2ce7c997e1971b7184ee37deb6ea9922ef5163c6ee5aa3c274b05f9e12fa \
+    --hash=sha256:c823bc3971c44ab93e611ab1a46b1eafeae474c0c844aff4b7474287b75fe49c \
+    --hash=sha256:ce0cdc074d540265bfeb31336e678b4e37316849d13b308607efa527e981f5c2 \
+    --hash=sha256:d1720b4f14c78a3089562b8875b53e36b51c97c51adc53325a69b79b4b48ebcb \
+    --hash=sha256:d183cf9c797a5291e8301790ed6d053480ed94070637bfaad914dd38b0981f67 \
+    --hash=sha256:d9010c31cd6fa59438da4e58a7f19e4753f7f264300cd152e7f90d4602449762 \
+    --hash=sha256:d9e5e4a85bdb56d224f412d9c98ae4cbd032cc4f3161818f692cd81766eee65a \
+    --hash=sha256:da1dee8948d2137bb51fbb8a53cce6b1bcc86003c6b42565f008438b806cccd8 \
+    --hash=sha256:df9270660711670e68803107d55c2b5949c2e0f2e4896da176e1ecfc068b974a \
+    --hash=sha256:e00e3505cd80440f6c98c6d69269dcc2a119f86ad0a9fd70bccc59504bebd68a \
+    --hash=sha256:e48d5021a84d341bcaf95c8460b152cfbad770d28e5fe14a768988c461b821bc \
+    --hash=sha256:e7f8b04d83483577fd9200461b057c9f14ced334dcb053090cea1da9c8321a91 \
+    --hash=sha256:edfe3341033a6b53a5c522c802deb2079eee5cbfbb0af032a55064bd65c73a23 \
+    --hash=sha256:ef9c33cc5cbca35808f6c74be11eb7f5f6b14d2311be84a15b594bd3e58b5527 \
+    --hash=sha256:f2d4324a98062be0525d16f768a03e0bbb3b9fe301ceee99611dc9a7953124e6 \
+    --hash=sha256:f3935f82f6f4a3820270842e90456ebad3af15810cf65932bd24da4463bc0a4c \
+    --hash=sha256:f614ab0c76397661b90b6851a030004dac502e48260ea10f2441abd2207fbcc7 \
+    --hash=sha256:f7db54c7914cc99d901d93a34704833568d86c20925b2762f9fa779f9cd2e70f \
+    --hash=sha256:fbc6264158392bad9df19537e872d476f7c57adf718944cc1e4495cbabf38e2a \
+    --hash=sha256:fe2fb38c2ed905a2582948e2de560675e9dfbee94c6d5ccdb1301c6d0a5bf092 \
+    --hash=sha256:ffe595f10566f8276b76dc3a11ae4bb7eba1aac8ddd75811736a15b0d5311414
+    # via gcsfs
+aiosignal==1.3.1 \
+    --hash=sha256:54cd96e15e1649b75d6c87526a6ff0b6c1b0dd3459f43d9ca11d48c339b68cfc \
+    --hash=sha256:f8376fb07dd1e86a584e4fcdec80b36b7f81aac666ebc724e2c090300dd83b17
+    # via aiohttp
+asciitree==0.3.3 \
+    --hash=sha256:4aa4b9b649f85e3fcb343363d97564aa1fb62e249677f2e18a96765145cc0f6e
+    # via zarr
+async-timeout==4.0.3 \
+    --hash=sha256:4640d96be84d82d02ed59ea2b7105a0f7b33abe8703703cd0ab0bf87c427522f \
+    --hash=sha256:7405140ff1230c310e51dc27b3145b9092d659ce68ff733fb0cefe3ee42be028
+    # via aiohttp
 attrs==24.2.0 \
     --hash=sha256:5cfb1b9148b5b086569baec03f20d7b6bf3bcacc9a42bebf87ffaaca362f6346 \
     --hash=sha256:81921eb96de3191c8258c199618104dd27ac608d9366f5e35d011eae1867ede2
-    # via e3nn-jax
+    # via
+    #   aiohttp
+    #   e3nn-jax
+cachetools==5.5.0 \
+    --hash=sha256:02134e8439cdc2ffb62023ce1debca2944c3f289d66bb17ead3ab3dede74b292 \
+    --hash=sha256:2cc24fb4cbe39633fb7badd9db9ca6295d766d9c2995f245725a46715d050f2a
+    # via google-auth
 certifi==2024.8.30 \
     --hash=sha256:922820b53db7a7257ffbda3f597266d435245903d80737e34f8a45ff3e3230d8 \
     --hash=sha256:bec941d2aa8195e248a60b31ff9f0558284cf01a52591ceda73ea9afffd69fd9
@@ -120,14 +235,78 @@ chex==0.1.87 \
     --hash=sha256:0096d89cc8d898bb521ef4bfbf5c24549022b0e5b301f529ab57238896fe6c5d \
     --hash=sha256:ce536475661fd96d21be0c1728ecdbedd03f8ff950c662dfc338c92ea782cb16
     # via optax
+cloudpickle==3.1.0 \
+    --hash=sha256:81a929b6e3c7335c863c771d673d105f02efdb89dfaba0c90495d1c64796601b \
+    --hash=sha256:fe11acda67f61aaaec473e3afe030feb131d78a43461b718185363384f1ba12e
+    # via tensorflow-probability
 contextlib2==21.6.0 \
     --hash=sha256:3fbdb64466afd23abaf6c977627b75b6139a5a3e8ce38405c5b413aed7a0471f \
     --hash=sha256:ab1e2bfe1d01d968e1b7e8d9023bc51ef3509bba217bb730cee3827e1ee82869
     # via ml-collections
+decorator==5.1.1 \
+    --hash=sha256:637996211036b6385ef91435e4fae22989472f9d571faba8927ba8253acbc330 \
+    --hash=sha256:b8c3f85900b9dc423225913c5aace94729fe1fa9763b38939a95226f02d37186
+    # via
+    #   gcsfs
+    #   tensorflow-probability
+dinosaur-dycore==1.0.0 \
+    --hash=sha256:0ea3807a83d0c5220ff7a1cd548c11d82091d455d17d8f7697b985b6283d2ef9 \
+    --hash=sha256:3505d59eb9d90ad58572f93fa87db5536b9df0663f95701d50c8ccb027e767e1
+    # via neuralgcm
 dm-haiku==0.0.12 \
     --hash=sha256:7448a43a6486bff95253f84e18eacc607d9c1256592573117a9d1d23e2780706 \
     --hash=sha256:ba0b3acf71433156737fe342c486da11727e5e6c9e054245f4f9b8f0b53eb608
-    # via jax-md
+    # via
+    #   jax-md
+    #   neuralgcm
+dm-tree==0.1.8 \
+    --hash=sha256:054b461f8176f4bce7a21f7b1870f873a1ced3bdbe1282c816c550bb43c71fa6 \
+    --hash=sha256:09964470f76a5201aff2e8f9b26842976de7889300676f927930f6285e256760 \
+    --hash=sha256:0d3172394079a86c3a759179c65f64c48d1a42b89495fcf38976d11cc3bb952c \
+    --hash=sha256:0e9620ccf06393eb6b613b5e366469304622d4ea96ae6540b28a33840e6c89cf \
+    --hash=sha256:0fcaabbb14e7980377439e7140bd05552739ca5e515ecb3119f234acee4b9430 \
+    --hash=sha256:1607ce49aa42f010d1e5e616d92ce899d66835d4d8bea49679582435285515de \
+    --hash=sha256:181c35521d480d0365f39300542cb6cd7fd2b77351bb43d7acfda15aef63b317 \
+    --hash=sha256:1d7c26e431fc93cc7e0cba867eb000db6a05f6f2b25af11ac4e9dada88fc5bca \
+    --hash=sha256:1fe962015b2fe1282892b28ebe962faed53c7f98d942da9a4625cbf27baef913 \
+    --hash=sha256:250b692fb75f45f02e2f58fbef9ab338904ef334b90557565621fa251df267cf \
+    --hash=sha256:2869228d9c619074de501a3c10dc7f07c75422f8fab36ecdcb859b6f1b1ec3ef \
+    --hash=sha256:28c52cbf4f8b3dbd0beaedf44f69fa85eec5e9dede612e08035e06ada6ec9426 \
+    --hash=sha256:2f7915660f59c09068e428613c480150180df1060561fd0d1470684ae7007bd1 \
+    --hash=sha256:343a4a4ebaa127451ff971254a4be4084eb4bdc0b2513c32b46f6f728fd03f9e \
+    --hash=sha256:35cc164a79336bfcfafb47e5f297898359123bbd3330c1967f0c4994f9cf9f60 \
+    --hash=sha256:378cc8ad93c5fe3590f405a309980721f021c790ca1bdf9b15bb1d59daec57f5 \
+    --hash=sha256:39070ba268c0491af9fe7a58644d99e8b4f2cde6e5884ba3380bddc84ed43d5f \
+    --hash=sha256:435227cf3c5dc63f4de054cf3d00183790bd9ead4c3623138c74dde7f67f521b \
+    --hash=sha256:5483dca4d7eb1a0d65fe86d3b6a53ae717face83c1f17e0887b1a4a64ae5c410 \
+    --hash=sha256:694c3654cfd2a81552c08ec66bb5c4a3d48fa292b9a181880fb081c36c5b9134 \
+    --hash=sha256:75c5d528bb992981c20793b6b453e91560784215dffb8a5440ba999753c14ceb \
+    --hash=sha256:803bfc53b4659f447ac694dbd04235f94a73ef7c1fd1e0df7c84ac41e0bc963b \
+    --hash=sha256:81fce77f22a302d7a5968aebdf4efafef4def7ce96528719a354e6990dcd49c7 \
+    --hash=sha256:83b7764de0d855338abefc6e3ee9fe40d301668310aa3baea3f778ff051f4393 \
+    --hash=sha256:8c60a7eadab64c2278861f56bca320b2720f163dca9d7558103c3b77f2416571 \
+    --hash=sha256:8ed3564abed97c806db122c2d3e1a2b64c74a63debe9903aad795167cc301368 \
+    --hash=sha256:94d3f0826311f45ee19b75f5b48c99466e4218a0489e81c0f0167bda50cacf22 \
+    --hash=sha256:96a548a406a6fb15fe58f6a30a57ff2f2aafbf25f05afab00c8f5e5977b6c715 \
+    --hash=sha256:a5d819c38c03f0bb5b3b3703c60e4b170355a0fc6b5819325bf3d4ceb3ae7e80 \
+    --hash=sha256:ad16ceba90a56ec47cf45b21856d14962ac314787975ef786efb5e6e9ca75ec7 \
+    --hash=sha256:af4b3d372f2477dcd89a6e717e4a575ca35ccc20cc4454a8a4b6f8838a00672d \
+    --hash=sha256:b095ba4f8ca1ba19350fd53cf1f8f3eb0bd406aa28af64a6dfc86707b32a810a \
+    --hash=sha256:b9bd9b9ccb59409d33d51d84b7668010c04c2af7d4a371632874c1ca356cff3d \
+    --hash=sha256:b9f89a454e98806b44fe9d40ec9eee61f848388f7e79ac2371a55679bd5a3ac6 \
+    --hash=sha256:bb2d109f42190225112da899b9f3d46d0d5f26aef501c61e43529fe9322530b5 \
+    --hash=sha256:c0a94aba18a35457a1b5cd716fd7b46c5dafdc4cf7869b4bae665b91c4682a8e \
+    --hash=sha256:c5c8c12e3fda754ef6af94161bacdaeda816d941995fac415d6855c6c386af68 \
+    --hash=sha256:d1612fcaecd79023dbc6a6ae48d51a80beb5c385d6f3f6d71688e57bc8d07de8 \
+    --hash=sha256:d16e1f2a073604cfcc09f7131ae8d534674f43c3aef4c25742eae295bc60d04f \
+    --hash=sha256:d20f2faa3672b52e5013f4077117bfb99c4cfc0b445d3bde1584c34032b57436 \
+    --hash=sha256:d40fa4106ca6edc66760246a08f500ec0c85ef55c762fb4a363f6ee739ba02ee \
+    --hash=sha256:de287fabc464b8734be251e46e06aa9aa1001f34198da2b6ce07bd197172b9cb \
+    --hash=sha256:e4d714371bb08839e4e5e29024fc95832d9affe129825ef38836b143028bd144 \
+    --hash=sha256:ea9e59e0451e7d29aece402d9f908f2e2a80922bcde2ebfd5dcb07750fcbfee8 \
+    --hash=sha256:f7ac31b9aecccb2c6e1ab29706f6ded3eba0c2c69c770322c9c685929c3d6afb \
+    --hash=sha256:fa42a605d099ee7d41ba2b5fb75e21423951fd26e5d50583a00471238fb3021d
+    # via tensorflow-probability
 e3nn-jax==0.20.7 \
     --hash=sha256:7764edb99d19465e574a23857e5f97420b4d720f06ba8bc18a01f9f2b8301d1c \
     --hash=sha256:8a1b6f0c00feed565358fceec4c916a623412b6328945fce272f0c33e52b4bfc
@@ -142,12 +321,24 @@ etils[epath,epy]==1.9.4 \
     # via
     #   optax
     #   orbax-checkpoint
+fasteners==0.19 \
+    --hash=sha256:758819cb5d94cdedf4e836988b74de396ceacb8e2794d21f82d131fd9ee77237 \
+    --hash=sha256:b4f37c3ac52d8a445af3a66bce57b33b5e90b97c696b7b984f530cf8f0ded09c
+    # via zarr
 flax==0.9.0 \
     --hash=sha256:12cd8f7162165ddd56877fb1cd9a4fcb47a31569e4c5343eeb59a36369fa2cfe \
     --hash=sha256:8b7f361eed0f5324e81f9dc8d02ea53da5f993d7c2e37e7aa5b37d3f6331dd53
     # via
     #   dm-haiku
     #   jax-md
+flexcache==0.3 \
+    --hash=sha256:18743bd5a0621bfe2cf8d519e4c3bfdf57a269c15d1ced3fb4b64e0ff4600656 \
+    --hash=sha256:d43c9fea82336af6e0115e308d9d33a185390b8346a017564611f1466dcd2e32
+    # via pint
+flexparser==0.4 \
+    --hash=sha256:266d98905595be2ccc5da964fe0a2c3526fbbffdc45b65b3146d75db992ef6b2 \
+    --hash=sha256:3738b456192dcb3e15620f324c447721023c0293f6af9955b481e91d00179846
+    # via pint
 frozendict==2.4.4 \
     --hash=sha256:07c3a5dee8bbb84cba770e273cdbf2c87c8e035903af8f781292d72583416801 \
     --hash=sha256:12a342e439aef28ccec533f0253ea53d75fe9102bd6ea928ff530e76eac38906 \
@@ -183,10 +374,187 @@ frozendict==2.4.4 \
     --hash=sha256:e1b941132d79ce72d562a13341d38fc217bc1ee24d8c35a20d754e79ff99e038 \
     --hash=sha256:fefeb700bc7eb8b4c2dc48704e4221860d254c8989fb53488540bc44e44a1ac2
     # via jax-md
+frozenlist==1.5.0 \
+    --hash=sha256:000a77d6034fbad9b6bb880f7ec073027908f1b40254b5d6f26210d2dab1240e \
+    --hash=sha256:03d33c2ddbc1816237a67f66336616416e2bbb6beb306e5f890f2eb22b959cdf \
+    --hash=sha256:04a5c6babd5e8fb7d3c871dc8b321166b80e41b637c31a995ed844a6139942b6 \
+    --hash=sha256:0996c66760924da6e88922756d99b47512a71cfd45215f3570bf1e0b694c206a \
+    --hash=sha256:0cc974cc93d32c42e7b0f6cf242a6bd941c57c61b618e78b6c0a96cb72788c1d \
+    --hash=sha256:0f253985bb515ecd89629db13cb58d702035ecd8cfbca7d7a7e29a0e6d39af5f \
+    --hash=sha256:11aabdd62b8b9c4b84081a3c246506d1cddd2dd93ff0ad53ede5defec7886b28 \
+    --hash=sha256:12f78f98c2f1c2429d42e6a485f433722b0061d5c0b0139efa64f396efb5886b \
+    --hash=sha256:140228863501b44b809fb39ec56b5d4071f4d0aa6d216c19cbb08b8c5a7eadb9 \
+    --hash=sha256:1431d60b36d15cda188ea222033eec8e0eab488f39a272461f2e6d9e1a8e63c2 \
+    --hash=sha256:15538c0cbf0e4fa11d1e3a71f823524b0c46299aed6e10ebb4c2089abd8c3bec \
+    --hash=sha256:15b731db116ab3aedec558573c1a5eec78822b32292fe4f2f0345b7f697745c2 \
+    --hash=sha256:17dcc32fc7bda7ce5875435003220a457bcfa34ab7924a49a1c19f55b6ee185c \
+    --hash=sha256:1893f948bf6681733aaccf36c5232c231e3b5166d607c5fa77773611df6dc336 \
+    --hash=sha256:189f03b53e64144f90990d29a27ec4f7997d91ed3d01b51fa39d2dbe77540fd4 \
+    --hash=sha256:1a8ea951bbb6cacd492e3948b8da8c502a3f814f5d20935aae74b5df2b19cf3d \
+    --hash=sha256:1b96af8c582b94d381a1c1f51ffaedeb77c821c690ea5f01da3d70a487dd0a9b \
+    --hash=sha256:1e76bfbc72353269c44e0bc2cfe171900fbf7f722ad74c9a7b638052afe6a00c \
+    --hash=sha256:2150cc6305a2c2ab33299453e2968611dacb970d2283a14955923062c8d00b10 \
+    --hash=sha256:226d72559fa19babe2ccd920273e767c96a49b9d3d38badd7c91a0fdeda8ea08 \
+    --hash=sha256:237f6b23ee0f44066219dae14c70ae38a63f0440ce6750f868ee08775073f942 \
+    --hash=sha256:29d94c256679247b33a3dc96cce0f93cbc69c23bf75ff715919332fdbb6a32b8 \
+    --hash=sha256:2b5e23253bb709ef57a8e95e6ae48daa9ac5f265637529e4ce6b003a37b2621f \
+    --hash=sha256:2d0da8bbec082bf6bf18345b180958775363588678f64998c2b7609e34719b10 \
+    --hash=sha256:2f3f7a0fbc219fb4455264cae4d9f01ad41ae6ee8524500f381de64ffaa077d5 \
+    --hash=sha256:30c72000fbcc35b129cb09956836c7d7abf78ab5416595e4857d1cae8d6251a6 \
+    --hash=sha256:31115ba75889723431aa9a4e77d5f398f5cf976eea3bdf61749731f62d4a4a21 \
+    --hash=sha256:31a9ac2b38ab9b5a8933b693db4939764ad3f299fcaa931a3e605bc3460e693c \
+    --hash=sha256:366d8f93e3edfe5a918c874702f78faac300209a4d5bf38352b2c1bdc07a766d \
+    --hash=sha256:374ca2dabdccad8e2a76d40b1d037f5bd16824933bf7bcea3e59c891fd4a0923 \
+    --hash=sha256:44c49271a937625619e862baacbd037a7ef86dd1ee215afc298a417ff3270608 \
+    --hash=sha256:45e0896250900b5aa25180f9aec243e84e92ac84bd4a74d9ad4138ef3f5c97de \
+    --hash=sha256:498524025a5b8ba81695761d78c8dd7382ac0b052f34e66939c42df860b8ff17 \
+    --hash=sha256:50cf5e7ee9b98f22bdecbabf3800ae78ddcc26e4a435515fc72d97903e8488e0 \
+    --hash=sha256:52ef692a4bc60a6dd57f507429636c2af8b6046db8b31b18dac02cbc8f507f7f \
+    --hash=sha256:561eb1c9579d495fddb6da8959fd2a1fca2c6d060d4113f5844b433fc02f2641 \
+    --hash=sha256:5a3ba5f9a0dfed20337d3e966dc359784c9f96503674c2faf015f7fe8e96798c \
+    --hash=sha256:5b6a66c18b5b9dd261ca98dffcb826a525334b2f29e7caa54e182255c5f6a65a \
+    --hash=sha256:5c28f4b5dbef8a0d8aad0d4de24d1e9e981728628afaf4ea0792f5d0939372f0 \
+    --hash=sha256:5d7f5a50342475962eb18b740f3beecc685a15b52c91f7d975257e13e029eca9 \
+    --hash=sha256:6321899477db90bdeb9299ac3627a6a53c7399c8cd58d25da094007402b039ab \
+    --hash=sha256:6482a5851f5d72767fbd0e507e80737f9c8646ae7fd303def99bfe813f76cf7f \
+    --hash=sha256:666534d15ba8f0fda3f53969117383d5dc021266b3c1a42c9ec4855e4b58b9d3 \
+    --hash=sha256:683173d371daad49cffb8309779e886e59c2f369430ad28fe715f66d08d4ab1a \
+    --hash=sha256:6e9080bb2fb195a046e5177f10d9d82b8a204c0736a97a153c2466127de87784 \
+    --hash=sha256:73f2e31ea8dd7df61a359b731716018c2be196e5bb3b74ddba107f694fbd7604 \
+    --hash=sha256:7437601c4d89d070eac8323f121fcf25f88674627505334654fd027b091db09d \
+    --hash=sha256:76e4753701248476e6286f2ef492af900ea67d9706a0155335a40ea21bf3b2f5 \
+    --hash=sha256:7707a25d6a77f5d27ea7dc7d1fc608aa0a478193823f88511ef5e6b8a48f9d03 \
+    --hash=sha256:7948140d9f8ece1745be806f2bfdf390127cf1a763b925c4a805c603df5e697e \
+    --hash=sha256:7a1a048f9215c90973402e26c01d1cff8a209e1f1b53f72b95c13db61b00f953 \
+    --hash=sha256:7d57d8f702221405a9d9b40f9da8ac2e4a1a8b5285aac6100f3393675f0a85ee \
+    --hash=sha256:7f3c8c1dacd037df16e85227bac13cca58c30da836c6f936ba1df0c05d046d8d \
+    --hash=sha256:81d5af29e61b9c8348e876d442253723928dce6433e0e76cd925cd83f1b4b817 \
+    --hash=sha256:828afae9f17e6de596825cf4228ff28fbdf6065974e5ac1410cecc22f699d2b3 \
+    --hash=sha256:87f724d055eb4785d9be84e9ebf0f24e392ddfad00b3fe036e43f489fafc9039 \
+    --hash=sha256:8969190d709e7c48ea386db202d708eb94bdb29207a1f269bab1196ce0dcca1f \
+    --hash=sha256:90646abbc7a5d5c7c19461d2e3eeb76eb0b204919e6ece342feb6032c9325ae9 \
+    --hash=sha256:91d6c171862df0a6c61479d9724f22efb6109111017c87567cfeb7b5d1449fdf \
+    --hash=sha256:9272fa73ca71266702c4c3e2d4a28553ea03418e591e377a03b8e3659d94fa76 \
+    --hash=sha256:92b5278ed9d50fe610185ecd23c55d8b307d75ca18e94c0e7de328089ac5dcba \
+    --hash=sha256:97160e245ea33d8609cd2b8fd997c850b56db147a304a262abc2b3be021a9171 \
+    --hash=sha256:977701c081c0241d0955c9586ffdd9ce44f7a7795df39b9151cd9a6fd0ce4cfb \
+    --hash=sha256:9b7dc0c4338e6b8b091e8faf0db3168a37101943e687f373dce00959583f7439 \
+    --hash=sha256:9b93d7aaa36c966fa42efcaf716e6b3900438632a626fb09c049f6a2f09fc631 \
+    --hash=sha256:9bbcdfaf4af7ce002694a4e10a0159d5a8d20056a12b05b45cea944a4953f972 \
+    --hash=sha256:9c2623347b933fcb9095841f1cc5d4ff0b278addd743e0e966cb3d460278840d \
+    --hash=sha256:a2fe128eb4edeabe11896cb6af88fca5346059f6c8d807e3b910069f39157869 \
+    --hash=sha256:a72b7a6e3cd2725eff67cd64c8f13335ee18fc3c7befc05aed043d24c7b9ccb9 \
+    --hash=sha256:a9fe0f1c29ba24ba6ff6abf688cb0b7cf1efab6b6aa6adc55441773c252f7411 \
+    --hash=sha256:b97f7b575ab4a8af9b7bc1d2ef7f29d3afee2226bd03ca3875c16451ad5a7723 \
+    --hash=sha256:bdac3c7d9b705d253b2ce370fde941836a5f8b3c5c2b8fd70940a3ea3af7f4f2 \
+    --hash=sha256:c03eff4a41bd4e38415cbed054bbaff4a075b093e2394b6915dca34a40d1e38b \
+    --hash=sha256:c16d2fa63e0800723139137d667e1056bee1a1cf7965153d2d104b62855e9b99 \
+    --hash=sha256:c1fac3e2ace2eb1052e9f7c7db480818371134410e1f5c55d65e8f3ac6d1407e \
+    --hash=sha256:ce3aa154c452d2467487765e3adc730a8c153af77ad84096bc19ce19a2400840 \
+    --hash=sha256:cee6798eaf8b1416ef6909b06f7dc04b60755206bddc599f52232606e18179d3 \
+    --hash=sha256:d1b3eb7b05ea246510b43a7e53ed1653e55c2121019a97e60cad7efb881a97bb \
+    --hash=sha256:d994863bba198a4a518b467bb971c56e1db3f180a25c6cf7bb1949c267f748c3 \
+    --hash=sha256:dd47a5181ce5fcb463b5d9e17ecfdb02b678cca31280639255ce9d0e5aa67af0 \
+    --hash=sha256:dd94994fc91a6177bfaafd7d9fd951bc8689b0a98168aa26b5f543868548d3ca \
+    --hash=sha256:de537c11e4aa01d37db0d403b57bd6f0546e71a82347a97c6a9f0dcc532b3a45 \
+    --hash=sha256:df6e2f325bfee1f49f81aaac97d2aa757c7646534a06f8f577ce184afe2f0a9e \
+    --hash=sha256:e66cc454f97053b79c2ab09c17fbe3c825ea6b4de20baf1be28919460dd7877f \
+    --hash=sha256:e79225373c317ff1e35f210dd5f1344ff31066ba8067c307ab60254cd3a78ad5 \
+    --hash=sha256:f1577515d35ed5649d52ab4319db757bb881ce3b2b796d7283e6634d99ace307 \
+    --hash=sha256:f1e6540b7fa044eee0bb5111ada694cf3dc15f2b0347ca125ee9ca984d5e9e6e \
+    --hash=sha256:f2ac49a9bedb996086057b75bf93538240538c6d9b38e57c82d51f75a73409d2 \
+    --hash=sha256:f47c9c9028f55a04ac254346e92977bf0f166c483c74b4232bee19a6697e4778 \
+    --hash=sha256:f5f9da7f5dbc00a604fe74aa02ae7c98bcede8a3b8b9666f9f86fc13993bc71a \
+    --hash=sha256:fd74520371c3c4175142d02a976aee0b4cb4a7cc912a60586ffd8d5929979b30 \
+    --hash=sha256:feeb64bc9bcc6b45c6311c9e9b99406660a9c05ca8a5b30d14a78555088b0b3a
+    # via
+    #   aiohttp
+    #   aiosignal
 fsspec==2024.9.0 \
     --hash=sha256:4b0afb90c2f21832df142f292649035d80b421f60a9e1c027802e5a0da2b04e8 \
     --hash=sha256:a0947d552d8a6efa72cc2c730b12c41d043509156966cca4fb157b0f2a0c574b
-    # via etils
+    # via
+    #   dinosaur-dycore
+    #   etils
+    #   gcsfs
+gast==0.6.0 \
+    --hash=sha256:52b182313f7330389f72b069ba00f174cfe2a06411099547288839c6cbafbd54 \
+    --hash=sha256:88fc5300d32c7ac6ca7b515310862f71e6fdf2c029bbec7c66c0f5dd47b6b1fb
+    # via tensorflow-probability
+gcsfs==2024.9.0.post1 \
+    --hash=sha256:7ca70ee9d7c7dbce1a3e36b4883e14102c2d7b4284f49e242843a437bc684684 \
+    --hash=sha256:f3ab9d3bedc45da8cf40baed7c3a1e1694e8f599160d9138d78f0ef25e4a3ca1
+    # via -r builddeps/test-requirements.txt
+gin-config==0.5.0 \
+    --hash=sha256:0c6ea5026ded927c8c93c990b01c695257c1df446e45e549a158cfbc79e19ed6 \
+    --hash=sha256:bddb7ca221ea2b46cdb59321e79fecf02d6e3b728906047fcd4076c297609fd6
+    # via neuralgcm
+google-api-core==2.22.0 \
+    --hash=sha256:26f8d76b96477db42b55fd02a33aae4a42ec8b86b98b94969b7333a2c828bf35 \
+    --hash=sha256:a6652b6bd51303902494998626653671703c420f6f4c88cfd3f50ed723e9d021
+    # via
+    #   google-cloud-core
+    #   google-cloud-storage
+google-auth==2.36.0 \
+    --hash=sha256:51a15d47028b66fd36e5c64a82d2d57480075bccc7da37cde257fc94177a61fb \
+    --hash=sha256:545e9618f2df0bcbb7dcbc45a546485b1212624716975a1ea5ae8149ce769ab1
+    # via
+    #   gcsfs
+    #   google-api-core
+    #   google-auth-oauthlib
+    #   google-cloud-core
+    #   google-cloud-storage
+google-auth-oauthlib==1.2.1 \
+    --hash=sha256:2d58a27262d55aa1b87678c3ba7142a080098cbc2024f903c62355deb235d91f \
+    --hash=sha256:afd0cad092a2eaa53cd8e8298557d6de1034c6cb4a740500b5357b648af97263
+    # via gcsfs
+google-cloud-core==2.4.1 \
+    --hash=sha256:9b7749272a812bde58fff28868d0c5e2f585b82f37e09a1f6ed2d4d10f134073 \
+    --hash=sha256:a9e6a4422b9ac5c29f79a0ede9485473338e2ce78d91f2370c01e730eab22e61
+    # via google-cloud-storage
+google-cloud-storage==2.18.2 \
+    --hash=sha256:97a4d45c368b7d401ed48c4fdfe86e1e1cb96401c9e199e419d289e2c0370166 \
+    --hash=sha256:aaf7acd70cdad9f274d29332673fcab98708d0e1f4dceb5a5356aaef06af4d99
+    # via gcsfs
+google-crc32c==1.6.0 \
+    --hash=sha256:05e2d8c9a2f853ff116db9706b4a27350587f341eda835f46db3c0a8c8ce2f24 \
+    --hash=sha256:18e311c64008f1f1379158158bb3f0c8d72635b9eb4f9545f8cf990c5668e59d \
+    --hash=sha256:236c87a46cdf06384f614e9092b82c05f81bd34b80248021f729396a78e55d7e \
+    --hash=sha256:35834855408429cecf495cac67ccbab802de269e948e27478b1e47dfb6465e57 \
+    --hash=sha256:386122eeaaa76951a8196310432c5b0ef3b53590ef4c317ec7588ec554fec5d2 \
+    --hash=sha256:40b05ab32a5067525670880eb5d169529089a26fe35dce8891127aeddc1950e8 \
+    --hash=sha256:48abd62ca76a2cbe034542ed1b6aee851b6f28aaca4e6551b5599b6f3ef175cc \
+    --hash=sha256:50cf2a96da226dcbff8671233ecf37bf6e95de98b2a2ebadbfdf455e6d05df42 \
+    --hash=sha256:51c4f54dd8c6dfeb58d1df5e4f7f97df8abf17a36626a217f169893d1d7f3e9f \
+    --hash=sha256:5bcc90b34df28a4b38653c36bb5ada35671ad105c99cfe915fb5bed7ad6924aa \
+    --hash=sha256:62f6d4a29fea082ac4a3c9be5e415218255cf11684ac6ef5488eea0c9132689b \
+    --hash=sha256:6eceb6ad197656a1ff49ebfbbfa870678c75be4344feb35ac1edf694309413dc \
+    --hash=sha256:7aec8e88a3583515f9e0957fe4f5f6d8d4997e36d0f61624e70469771584c760 \
+    --hash=sha256:91ca8145b060679ec9176e6de4f89b07363d6805bd4760631ef254905503598d \
+    --hash=sha256:a184243544811e4a50d345838a883733461e67578959ac59964e43cca2c791e7 \
+    --hash=sha256:a9e4b426c3702f3cd23b933436487eb34e01e00327fac20c9aebb68ccf34117d \
+    --hash=sha256:bb0966e1c50d0ef5bc743312cc730b533491d60585a9a08f897274e57c3f70e0 \
+    --hash=sha256:bb8b3c75bd157010459b15222c3fd30577042a7060e29d42dabce449c087f2b3 \
+    --hash=sha256:bd5e7d2445d1a958c266bfa5d04c39932dc54093fa391736dbfdb0f1929c1fb3 \
+    --hash=sha256:c87d98c7c4a69066fd31701c4e10d178a648c2cac3452e62c6b24dc51f9fcc00 \
+    --hash=sha256:d2952396dc604544ea7476b33fe87faedc24d666fb0c2d5ac971a2b9576ab871 \
+    --hash=sha256:d8797406499f28b5ef791f339594b0b5fdedf54e203b5066675c406ba69d705c \
+    --hash=sha256:d9e9913f7bd69e093b81da4535ce27af842e7bf371cde42d1ae9e9bd382dc0e9 \
+    --hash=sha256:e2806553238cd076f0a55bddab37a532b53580e699ed8e5606d0de1f856b5205 \
+    --hash=sha256:ebab974b1687509e5c973b5c4b8b146683e101e102e17a86bd196ecaa4d099fc \
+    --hash=sha256:ed767bf4ba90104c1216b68111613f0d5926fb3780660ea1198fc469af410e9d \
+    --hash=sha256:f7a1fc29803712f80879b0806cb83ab24ce62fc8daf0569f2204a0cfd7f68ed4
+    # via
+    #   google-cloud-storage
+    #   google-resumable-media
+google-resumable-media==2.7.2 \
+    --hash=sha256:3ce7551e9fe6d99e9a126101d2536612bb73486721951e9562fee0f90c6ababa \
+    --hash=sha256:5280aed4629f2b60b847b0d42f9857fd4935c11af266744df33d8074cae92fe0
+    # via google-cloud-storage
+googleapis-common-protos==1.65.0 \
+    --hash=sha256:2972e6c496f435b92590fd54045060867f3fe9be2c82ab148fc8885035479a63 \
+    --hash=sha256:334a29d07cddc3aa01dee4988f9afd9b2916ee2ff49d6b757155dc0d197852c0
+    # via google-api-core
 humanize==4.11.0 \
     --hash=sha256:b53caaec8532bcb2fff70c8826f904c35943f8cecaca29d272d9df38092736c0 \
     --hash=sha256:e66f36020a2d5a974c504bd2555cf770621dbdbb6d82f94a6857c0b1ea2608be
@@ -194,7 +562,9 @@ humanize==4.11.0 \
 idna==3.10 \
     --hash=sha256:12f65c9b470abda6dc35cf8e63cc574b1c52b11df2c86030af0ac09b01b13ea9 \
     --hash=sha256:946d195a0d259cbba61165e88e65941f16e9b36ea6ddb97f00452bae8b1287d3
-    # via requests
+    # via
+    #   requests
+    #   yarl
 importlib-resources==6.4.5 \
     --hash=sha256:980862a1d16c9e147a59603677fa2aa5fd82b87f223b6cb870695bcfce830065 \
     --hash=sha256:ac29d5f956f01d5e4bb63102a5a19957f1b9175e45649977264a1416783bb717
@@ -206,12 +576,15 @@ jax==0.4.31 \
     #   -r builddeps/requirements.in
     #   -r builddeps/test-requirements.txt
     #   chex
+    #   dinosaur-dycore
     #   e3nn-jax
     #   flax
     #   jax-md
     #   jraph
+    #   neuralgcm
     #   optax
     #   orbax-checkpoint
+    #   tree-math
 jax-cuda12-pjrt==0.4.31 \
     --hash=sha256:3e77d1cfebeca06517254eb568f082037e1a2aa3ed8f63c543492ad8ab5a1585 \
     --hash=sha256:8961abb381d893a3c2392ad76ab2067a81f8f2514f3b47d2da3ac24283293fe0
@@ -247,19 +620,24 @@ jaxlib==0.4.31 \
     #   -r builddeps/requirements.in
     #   -r builddeps/test-requirements.txt
     #   chex
+    #   dinosaur-dycore
     #   e3nn-jax
     #   jax
     #   jax-md
     #   jraph
+    #   neuralgcm
     #   optax
     #   orbax-checkpoint
 jmp==0.0.4 \
     --hash=sha256:5dfeb0fd7c7a9f72a70fff0aab9d0cbfae32a809c02f4037ff3485ceb33e1730 \
     --hash=sha256:6aa7adbddf2bd574b28c7faf6e81a735eb11f53386447896909c6968dc36807d
     # via dm-haiku
+joblib==1.4.2 \
+    --hash=sha256:06d478d5674cbc267e7496a410ee875abd68e4340feff4490bcb7afb88060ae6 \
+    --hash=sha256:2382c5816b2636fbd20a09e0f4e9dad4736765fdfb7dca582943b9c1366b3f0e
+    # via scikit-learn
 jraph @ https://github.com/wsmoses/jraph/archive/b00d9a03db76c69a258a86df81638b9a2f28829f.tar.gz \
-    --hash=sha256:350fe37bf717f934f1f84fd3370a480b3178bfcb61dfa217c738971308c57625 \
-    --hash=sha256:c3ac3a0b224b344eb6d367e8bc312d95ea41bf825d01ea31b80dd8c22c0dd8b8
+    --hash=sha256:d996f47fff8c6ac8f5572e4220e3193f2542e9d3d6bc1defc16ac22dadd74dd4
     # via -r builddeps/test-requirements.txt
 markdown-it-py==3.0.0 \
     --hash=sha256:355216845c60bd96232cd8d8c40e8f9765cc86f46880e43a8fd22dc1a1a8cab1 \
@@ -366,10 +744,129 @@ msgpack==1.1.0 \
     # via
     #   flax
     #   orbax-checkpoint
+multidict==6.1.0 \
+    --hash=sha256:052e10d2d37810b99cc170b785945421141bf7bb7d2f8799d431e7db229c385f \
+    --hash=sha256:06809f4f0f7ab7ea2cabf9caca7d79c22c0758b58a71f9d32943ae13c7ace056 \
+    --hash=sha256:071120490b47aa997cca00666923a83f02c7fbb44f71cf7f136df753f7fa8761 \
+    --hash=sha256:0c3f390dc53279cbc8ba976e5f8035eab997829066756d811616b652b00a23a3 \
+    --hash=sha256:0e2b90b43e696f25c62656389d32236e049568b39320e2735d51f08fd362761b \
+    --hash=sha256:0e5f362e895bc5b9e67fe6e4ded2492d8124bdf817827f33c5b46c2fe3ffaca6 \
+    --hash=sha256:10524ebd769727ac77ef2278390fb0068d83f3acb7773792a5080f2b0abf7748 \
+    --hash=sha256:10a9b09aba0c5b48c53761b7c720aaaf7cf236d5fe394cd399c7ba662d5f9966 \
+    --hash=sha256:16e5f4bf4e603eb1fdd5d8180f1a25f30056f22e55ce51fb3d6ad4ab29f7d96f \
+    --hash=sha256:188215fc0aafb8e03341995e7c4797860181562380f81ed0a87ff455b70bf1f1 \
+    --hash=sha256:189f652a87e876098bbc67b4da1049afb5f5dfbaa310dd67c594b01c10388db6 \
+    --hash=sha256:1ca0083e80e791cffc6efce7660ad24af66c8d4079d2a750b29001b53ff59ada \
+    --hash=sha256:1e16bf3e5fc9f44632affb159d30a437bfe286ce9e02754759be5536b169b305 \
+    --hash=sha256:2090f6a85cafc5b2db085124d752757c9d251548cedabe9bd31afe6363e0aff2 \
+    --hash=sha256:20b9b5fbe0b88d0bdef2012ef7dee867f874b72528cf1d08f1d59b0e3850129d \
+    --hash=sha256:22ae2ebf9b0c69d206c003e2f6a914ea33f0a932d4aa16f236afc049d9958f4a \
+    --hash=sha256:22f3105d4fb15c8f57ff3959a58fcab6ce36814486500cd7485651230ad4d4ef \
+    --hash=sha256:23bfd518810af7de1116313ebd9092cb9aa629beb12f6ed631ad53356ed6b86c \
+    --hash=sha256:27e5fc84ccef8dfaabb09d82b7d179c7cf1a3fbc8a966f8274fcb4ab2eb4cadb \
+    --hash=sha256:3380252550e372e8511d49481bd836264c009adb826b23fefcc5dd3c69692f60 \
+    --hash=sha256:3702ea6872c5a2a4eeefa6ffd36b042e9773f05b1f37ae3ef7264b1163c2dcf6 \
+    --hash=sha256:37bb93b2178e02b7b618893990941900fd25b6b9ac0fa49931a40aecdf083fe4 \
+    --hash=sha256:3914f5aaa0f36d5d60e8ece6a308ee1c9784cd75ec8151062614657a114c4478 \
+    --hash=sha256:3a37ffb35399029b45c6cc33640a92bef403c9fd388acce75cdc88f58bd19a81 \
+    --hash=sha256:3c8b88a2ccf5493b6c8da9076fb151ba106960a2df90c2633f342f120751a9e7 \
+    --hash=sha256:3e97b5e938051226dc025ec80980c285b053ffb1e25a3db2a3aa3bc046bf7f56 \
+    --hash=sha256:3ec660d19bbc671e3a6443325f07263be452c453ac9e512f5eb935e7d4ac28b3 \
+    --hash=sha256:3efe2c2cb5763f2f1b275ad2bf7a287d3f7ebbef35648a9726e3b69284a4f3d6 \
+    --hash=sha256:483a6aea59cb89904e1ceabd2b47368b5600fb7de78a6e4a2c2987b2d256cf30 \
+    --hash=sha256:4867cafcbc6585e4b678876c489b9273b13e9fff9f6d6d66add5e15d11d926cb \
+    --hash=sha256:48e171e52d1c4d33888e529b999e5900356b9ae588c2f09a52dcefb158b27506 \
+    --hash=sha256:4a9cb68166a34117d6646c0023c7b759bf197bee5ad4272f420a0141d7eb03a0 \
+    --hash=sha256:4b820514bfc0b98a30e3d85462084779900347e4d49267f747ff54060cc33925 \
+    --hash=sha256:4e18b656c5e844539d506a0a06432274d7bd52a7487e6828c63a63d69185626c \
+    --hash=sha256:4e9f48f58c2c523d5a06faea47866cd35b32655c46b443f163d08c6d0ddb17d6 \
+    --hash=sha256:50b3a2710631848991d0bf7de077502e8994c804bb805aeb2925a981de58ec2e \
+    --hash=sha256:55b6d90641869892caa9ca42ff913f7ff1c5ece06474fbd32fb2cf6834726c95 \
+    --hash=sha256:57feec87371dbb3520da6192213c7d6fc892d5589a93db548331954de8248fd2 \
+    --hash=sha256:58130ecf8f7b8112cdb841486404f1282b9c86ccb30d3519faf301b2e5659133 \
+    --hash=sha256:5845c1fd4866bb5dd3125d89b90e57ed3138241540897de748cdf19de8a2fca2 \
+    --hash=sha256:59bfeae4b25ec05b34f1956eaa1cb38032282cd4dfabc5056d0a1ec4d696d3aa \
+    --hash=sha256:5b48204e8d955c47c55b72779802b219a39acc3ee3d0116d5080c388970b76e3 \
+    --hash=sha256:5c09fcfdccdd0b57867577b719c69e347a436b86cd83747f179dbf0cc0d4c1f3 \
+    --hash=sha256:6180c0ae073bddeb5a97a38c03f30c233e0a4d39cd86166251617d1bbd0af436 \
+    --hash=sha256:682b987361e5fd7a139ed565e30d81fd81e9629acc7d925a205366877d8c8657 \
+    --hash=sha256:6b5d83030255983181005e6cfbac1617ce9746b219bc2aad52201ad121226581 \
+    --hash=sha256:6bb5992037f7a9eff7991ebe4273ea7f51f1c1c511e6a2ce511d0e7bdb754492 \
+    --hash=sha256:73eae06aa53af2ea5270cc066dcaf02cc60d2994bbb2c4ef5764949257d10f43 \
+    --hash=sha256:76f364861c3bfc98cbbcbd402d83454ed9e01a5224bb3a28bf70002a230f73e2 \
+    --hash=sha256:820c661588bd01a0aa62a1283f20d2be4281b086f80dad9e955e690c75fb54a2 \
+    --hash=sha256:82176036e65644a6cc5bd619f65f6f19781e8ec2e5330f51aa9ada7504cc1926 \
+    --hash=sha256:87701f25a2352e5bf7454caa64757642734da9f6b11384c1f9d1a8e699758057 \
+    --hash=sha256:9079dfc6a70abe341f521f78405b8949f96db48da98aeb43f9907f342f627cdc \
+    --hash=sha256:90f8717cb649eea3504091e640a1b8568faad18bd4b9fcd692853a04475a4b80 \
+    --hash=sha256:957cf8e4b6e123a9eea554fa7ebc85674674b713551de587eb318a2df3e00255 \
+    --hash=sha256:99f826cbf970077383d7de805c0681799491cb939c25450b9b5b3ced03ca99f1 \
+    --hash=sha256:9f636b730f7e8cb19feb87094949ba54ee5357440b9658b2a32a5ce4bce53972 \
+    --hash=sha256:a114d03b938376557927ab23f1e950827c3b893ccb94b62fd95d430fd0e5cf53 \
+    --hash=sha256:a185f876e69897a6f3325c3f19f26a297fa058c5e456bfcff8015e9a27e83ae1 \
+    --hash=sha256:a7a9541cd308eed5e30318430a9c74d2132e9a8cb46b901326272d780bf2d423 \
+    --hash=sha256:aa466da5b15ccea564bdab9c89175c762bc12825f4659c11227f515cee76fa4a \
+    --hash=sha256:aaed8b0562be4a0876ee3b6946f6869b7bcdb571a5d1496683505944e268b160 \
+    --hash=sha256:ab7c4ceb38d91570a650dba194e1ca87c2b543488fe9309b4212694174fd539c \
+    --hash=sha256:ac10f4c2b9e770c4e393876e35a7046879d195cd123b4f116d299d442b335bcd \
+    --hash=sha256:b04772ed465fa3cc947db808fa306d79b43e896beb677a56fb2347ca1a49c1fa \
+    --hash=sha256:b1c416351ee6271b2f49b56ad7f308072f6f44b37118d69c2cad94f3fa8a40d5 \
+    --hash=sha256:b225d95519a5bf73860323e633a664b0d85ad3d5bede6d30d95b35d4dfe8805b \
+    --hash=sha256:b2f59caeaf7632cc633b5cf6fc449372b83bbdf0da4ae04d5be36118e46cc0aa \
+    --hash=sha256:b58c621844d55e71c1b7f7c498ce5aa6985d743a1a59034c57a905b3f153c1ef \
+    --hash=sha256:bf6bea52ec97e95560af5ae576bdac3aa3aae0b6758c6efa115236d9e07dae44 \
+    --hash=sha256:c08be4f460903e5a9d0f76818db3250f12e9c344e79314d1d570fc69d7f4eae4 \
+    --hash=sha256:c7053d3b0353a8b9de430a4f4b4268ac9a4fb3481af37dfe49825bf45ca24156 \
+    --hash=sha256:c943a53e9186688b45b323602298ab727d8865d8c9ee0b17f8d62d14b56f0753 \
+    --hash=sha256:ce2186a7df133a9c895dea3331ddc5ddad42cdd0d1ea2f0a51e5d161e4762f28 \
+    --hash=sha256:d093be959277cb7dee84b801eb1af388b6ad3ca6a6b6bf1ed7585895789d027d \
+    --hash=sha256:d094ddec350a2fb899fec68d8353c78233debde9b7d8b4beeafa70825f1c281a \
+    --hash=sha256:d1a9dd711d0877a1ece3d2e4fea11a8e75741ca21954c919406b44e7cf971304 \
+    --hash=sha256:d569388c381b24671589335a3be6e1d45546c2988c2ebe30fdcada8457a31008 \
+    --hash=sha256:d618649d4e70ac6efcbba75be98b26ef5078faad23592f9b51ca492953012429 \
+    --hash=sha256:d83a047959d38a7ff552ff94be767b7fd79b831ad1cd9920662db05fec24fe72 \
+    --hash=sha256:d8fff389528cad1618fb4b26b95550327495462cd745d879a8c7c2115248e399 \
+    --hash=sha256:da1758c76f50c39a2efd5e9859ce7d776317eb1dd34317c8152ac9251fc574a3 \
+    --hash=sha256:db7457bac39421addd0c8449933ac32d8042aae84a14911a757ae6ca3eef1392 \
+    --hash=sha256:e27bbb6d14416713a8bd7aaa1313c0fc8d44ee48d74497a0ff4c3a1b6ccb5167 \
+    --hash=sha256:e617fb6b0b6953fffd762669610c1c4ffd05632c138d61ac7e14ad187870669c \
+    --hash=sha256:e9aa71e15d9d9beaad2c6b9319edcdc0a49a43ef5c0a4c8265ca9ee7d6c67774 \
+    --hash=sha256:ec2abea24d98246b94913b76a125e855eb5c434f7c46546046372fe60f666351 \
+    --hash=sha256:f179dee3b863ab1c59580ff60f9d99f632f34ccb38bf67a33ec6b3ecadd0fd76 \
+    --hash=sha256:f4c035da3f544b1882bac24115f3e2e8760f10a0107614fc9839fd232200b875 \
+    --hash=sha256:f67f217af4b1ff66c68a87318012de788dd95fcfeb24cc889011f4e1c7454dfd \
+    --hash=sha256:f90c822a402cb865e396a504f9fc8173ef34212a342d92e362ca498cad308e28 \
+    --hash=sha256:ff3827aef427c89a25cc96ded1759271a93603aba9fb977a6d264648ebf989db
+    # via
+    #   aiohttp
+    #   yarl
 nest-asyncio==1.6.0 \
     --hash=sha256:6f172d5449aca15afd6c646851f4e31e02c598d553a667e38cafa997cfec55fe \
     --hash=sha256:87af6efd6b5e897c81050477ef65c62e2b2f35d51703cae01aff2905b1852e1c
     # via orbax-checkpoint
+neuralgcm==1.0.0 \
+    --hash=sha256:28f42b73188df92e35d5e187ef9ebd2b2a9922cfc46e05e0ae5202f451ef7c18 \
+    --hash=sha256:e75889894e8948211cec9bf18c89df5e223ef5a6d9a6444bb42f4b36143a5c27
+    # via -r builddeps/test-requirements.txt
+numcodecs==0.13.1 \
+    --hash=sha256:233bc7f26abce24d57e44ea8ebeb5cd17084690b4e7409dd470fdb75528d615f \
+    --hash=sha256:237b7171609e868a20fd313748494444458ccd696062f67e198f7f8f52000c15 \
+    --hash=sha256:2a86f5367af9168e30f99727ff03b27d849c31ad4522060dde0bce2923b3a8bc \
+    --hash=sha256:2eda97dd2f90add98df6d295f2c6ae846043396e3d51a739ca5db6c03b5eb666 \
+    --hash=sha256:3501a848adaddce98a71a262fee15cd3618312692aa419da77acd18af4a6a3f6 \
+    --hash=sha256:3f593c7506b0ab248961a3b13cb148cc6e8355662ff124ac591822310bc55ecf \
+    --hash=sha256:5195bea384a6428f8afcece793860b1ab0ae28143c853f0b2b20d55a8947c917 \
+    --hash=sha256:796b3e6740107e4fa624cc636248a1580138b3f1c579160f260f76ff13a4261b \
+    --hash=sha256:7a60d75179fd6692e301ddfb3b266d51eb598606dcae7b9fc57f986e8d65cb43 \
+    --hash=sha256:80d3071465f03522e776a31045ddf2cfee7f52df468b977ed3afdd7fe5869701 \
+    --hash=sha256:90d3065ae74c9342048ae0046006f99dcb1388b7288da5a19b3bddf9c30c3176 \
+    --hash=sha256:96add4f783c5ce57cc7e650b6cac79dd101daf887c479a00a29bc1487ced180b \
+    --hash=sha256:96e42f73c31b8c24259c5fac6adba0c3ebf95536e37749dc6c62ade2989dca28 \
+    --hash=sha256:a3cf37881df0898f3a9c0d4477df88133fe85185bffe57ba31bcc2fa207709bc \
+    --hash=sha256:da2230484e6102e5fa3cc1a5dd37ca1f92dfbd183d91662074d6f7574e3e8f53 \
+    --hash=sha256:e5db4824ebd5389ea30e54bc8aeccb82d514d28b6b68da6c536b8fa4596f4bca \
+    --hash=sha256:eda7d7823c9282e65234731fd6bd3986b1f9e035755f7fed248d7d366bb291ab
+    # via zarr
 numpy==2.1.0 \
     --hash=sha256:08801848a40aea24ce16c2ecde3b756f9ad756586fb2d13210939eb69b023f5b \
     --hash=sha256:0937e54c09f7a9a68da6889362ddd2ff584c02d015ec92672c099b61555f8911 \
@@ -426,6 +923,7 @@ numpy==2.1.0 \
     # via
     #   -r builddeps/test-requirements.txt
     #   chex
+    #   dinosaur-dycore
     #   dm-haiku
     #   e3nn-jax
     #   jax
@@ -434,11 +932,19 @@ numpy==2.1.0 \
     #   jmp
     #   jraph
     #   ml-dtypes
+    #   neuralgcm
+    #   numcodecs
     #   opt-einsum
     #   optax
     #   orbax-checkpoint
+    #   pandas
+    #   scikit-learn
     #   scipy
+    #   tensorflow-probability
     #   tensorstore
+    #   xarray
+    #   xarray-tensorstore
+    #   zarr
 nvidia-cublas-cu12==12.6.3.3 \
     --hash=sha256:e1f70bee38b964eac1907293b336bceb24498a4243e61eaf91a52977c59aebc4 \
     --hash=sha256:e531199ca4f1f764fb45bc1dde49a006f6765033f9c89c737e4553b9502ca1f5 \
@@ -497,6 +1003,10 @@ nvidia-nvjitlink-cu12==12.6.77 \
     #   nvidia-cufft-cu12
     #   nvidia-cusolver-cu12
     #   nvidia-cusparse-cu12
+oauthlib==3.2.2 \
+    --hash=sha256:8139f29aac13e25d502680e9e19963e83f16838d48a0d71c287fe40e7067fbca \
+    --hash=sha256:9859c40929662bec5d64f34d01c99e093149682a3f38915dc0655d5a633dd918
+    # via requests-oauthlib
 opt-einsum==3.3.0 \
     --hash=sha256:2455e59e3947d3c275477df7f5205b30635e266fe6dc300e3d9f9646bfcea147 \
     --hash=sha256:59f6475f77bbc37dcf7cd748519c0ec60722e91e63ca114e68821c0c54a46549
@@ -507,10 +1017,174 @@ optax==0.2.3 \
     # via
     #   flax
     #   jax-md
+    #   neuralgcm
 orbax-checkpoint==0.6.4 \
     --hash=sha256:366b4d528a7322e1b3d9ddcaed45c8515add0d2fc69c8975c30d98638543240f \
     --hash=sha256:b4f2608ee4d1da67f7619fc35ff9c928ecdf4ccf7546eeb43ecf38c2608b6dea
     # via flax
+packaging==24.1 \
+    --hash=sha256:026ed72c8ed3fcce5bf8950572258698927fd1dbda10a5e981cdf0ac37f4f002 \
+    --hash=sha256:5b8f2217dbdbd2f7f384c41c628544e6d52f2d0f53c6d0c3ea61aa5d1d7ff124
+    # via xarray
+pandas==2.2.3 \
+    --hash=sha256:062309c1b9ea12a50e8ce661145c6aab431b1e99530d3cd60640e255778bd43a \
+    --hash=sha256:15c0e1e02e93116177d29ff83e8b1619c93ddc9c49083f237d4312337a61165d \
+    --hash=sha256:1948ddde24197a0f7add2bdc4ca83bf2b1ef84a1bc8ccffd95eda17fd836ecb5 \
+    --hash=sha256:1db71525a1538b30142094edb9adc10be3f3e176748cd7acc2240c2f2e5aa3a4 \
+    --hash=sha256:22a9d949bfc9a502d320aa04e5d02feab689d61da4e7764b62c30b991c42c5f0 \
+    --hash=sha256:29401dbfa9ad77319367d36940cd8a0b3a11aba16063e39632d98b0e931ddf32 \
+    --hash=sha256:31d0ced62d4ea3e231a9f228366919a5ea0b07440d9d4dac345376fd8e1477ea \
+    --hash=sha256:3508d914817e153ad359d7e069d752cdd736a247c322d932eb89e6bc84217f28 \
+    --hash=sha256:37e0aced3e8f539eccf2e099f65cdb9c8aa85109b0be6e93e2baff94264bdc6f \
+    --hash=sha256:381175499d3802cde0eabbaf6324cce0c4f5d52ca6f8c377c29ad442f50f6348 \
+    --hash=sha256:38cf8125c40dae9d5acc10fa66af8ea6fdf760b2714ee482ca691fc66e6fcb18 \
+    --hash=sha256:3b71f27954685ee685317063bf13c7709a7ba74fc996b84fc6821c59b0f06468 \
+    --hash=sha256:3fc6873a41186404dad67245896a6e440baacc92f5b716ccd1bc9ed2995ab2c5 \
+    --hash=sha256:4850ba03528b6dd51d6c5d273c46f183f39a9baf3f0143e566b89450965b105e \
+    --hash=sha256:4f18ba62b61d7e192368b84517265a99b4d7ee8912f8708660fb4a366cc82667 \
+    --hash=sha256:56534ce0746a58afaf7942ba4863e0ef81c9c50d3f0ae93e9497d6a41a057645 \
+    --hash=sha256:59ef3764d0fe818125a5097d2ae867ca3fa64df032331b7e0917cf5d7bf66b13 \
+    --hash=sha256:5dbca4c1acd72e8eeef4753eeca07de9b1db4f398669d5994086f788a5d7cc30 \
+    --hash=sha256:5de54125a92bb4d1c051c0659e6fcb75256bf799a732a87184e5ea503965bce3 \
+    --hash=sha256:61c5ad4043f791b61dd4752191d9f07f0ae412515d59ba8f005832a532f8736d \
+    --hash=sha256:6374c452ff3ec675a8f46fd9ab25c4ad0ba590b71cf0656f8b6daa5202bca3fb \
+    --hash=sha256:63cc132e40a2e084cf01adf0775b15ac515ba905d7dcca47e9a251819c575ef3 \
+    --hash=sha256:66108071e1b935240e74525006034333f98bcdb87ea116de573a6a0dccb6c039 \
+    --hash=sha256:6dfcb5ee8d4d50c06a51c2fffa6cff6272098ad6540aed1a76d15fb9318194d8 \
+    --hash=sha256:7c2875855b0ff77b2a64a0365e24455d9990730d6431b9e0ee18ad8acee13dbd \
+    --hash=sha256:7eee9e7cea6adf3e3d24e304ac6b8300646e2a5d1cd3a3c2abed9101b0846761 \
+    --hash=sha256:800250ecdadb6d9c78eae4990da62743b857b470883fa27f652db8bdde7f6659 \
+    --hash=sha256:86976a1c5b25ae3f8ccae3a5306e443569ee3c3faf444dfd0f41cda24667ad57 \
+    --hash=sha256:8cd6d7cc958a3910f934ea8dbdf17b2364827bb4dafc38ce6eef6bb3d65ff09c \
+    --hash=sha256:99df71520d25fade9db7c1076ac94eb994f4d2673ef2aa2e86ee039b6746d20c \
+    --hash=sha256:a5a1595fe639f5988ba6a8e5bc9649af3baf26df3998a0abe56c02609392e0a4 \
+    --hash=sha256:ad5b65698ab28ed8d7f18790a0dc58005c7629f227be9ecc1072aa74c0c1d43a \
+    --hash=sha256:b1d432e8d08679a40e2a6d8b2f9770a5c21793a6f9f47fdd52c5ce1948a5a8a9 \
+    --hash=sha256:b8661b0238a69d7aafe156b7fa86c44b881387509653fdf857bebc5e4008ad42 \
+    --hash=sha256:ba96630bc17c875161df3818780af30e43be9b166ce51c9a18c1feae342906c2 \
+    --hash=sha256:bc6b93f9b966093cb0fd62ff1a7e4c09e6d546ad7c1de191767baffc57628f39 \
+    --hash=sha256:c124333816c3a9b03fbeef3a9f230ba9a737e9e5bb4060aa2107a86cc0a497fc \
+    --hash=sha256:cd8d0c3be0515c12fed0bdbae072551c8b54b7192c7b1fda0ba56059a0179698 \
+    --hash=sha256:d9c45366def9a3dd85a6454c0e7908f2b3b8e9c138f5dc38fed7ce720d8453ed \
+    --hash=sha256:f00d1345d84d8c86a63e476bb4955e46458b304b9575dcf71102b5c705320015 \
+    --hash=sha256:f3a255b2c19987fbbe62a9dfd6cff7ff2aa9ccab3fc75218fd4b7530f01efa24 \
+    --hash=sha256:fffb8ae78d8af97f849404f21411c95062db1496aeb3e56f146f0355c9989319
+    # via
+    #   dinosaur-dycore
+    #   neuralgcm
+    #   xarray
+pint==0.24.4 \
+    --hash=sha256:35275439b574837a6cd3020a5a4a73645eb125ce4152a73a2f126bf164b91b80 \
+    --hash=sha256:aa54926c8772159fcf65f82cc0d34de6768c151b32ad1deb0331291c38fe7659
+    # via dinosaur-dycore
+platformdirs==4.3.6 \
+    --hash=sha256:357fb2acbc885b0419afd3ce3ed34564c13c9b95c89360cd9563f73aa5e2b907 \
+    --hash=sha256:73e575e1408ab8103900836b97580d5307456908a03e92031bab39e4554cc3fb
+    # via pint
+propcache==0.2.0 \
+    --hash=sha256:00181262b17e517df2cd85656fcd6b4e70946fe62cd625b9d74ac9977b64d8d9 \
+    --hash=sha256:0e53cb83fdd61cbd67202735e6a6687a7b491c8742dfc39c9e01e80354956763 \
+    --hash=sha256:1235c01ddaa80da8235741e80815ce381c5267f96cc49b1477fdcf8c047ef325 \
+    --hash=sha256:140fbf08ab3588b3468932974a9331aff43c0ab8a2ec2c608b6d7d1756dbb6cb \
+    --hash=sha256:191db28dc6dcd29d1a3e063c3be0b40688ed76434622c53a284e5427565bbd9b \
+    --hash=sha256:1e41d67757ff4fbc8ef2af99b338bfb955010444b92929e9e55a6d4dcc3c4f09 \
+    --hash=sha256:1ec43d76b9677637a89d6ab86e1fef70d739217fefa208c65352ecf0282be957 \
+    --hash=sha256:20a617c776f520c3875cf4511e0d1db847a076d720714ae35ffe0df3e440be68 \
+    --hash=sha256:218db2a3c297a3768c11a34812e63b3ac1c3234c3a086def9c0fee50d35add1f \
+    --hash=sha256:22aa8f2272d81d9317ff5756bb108021a056805ce63dd3630e27d042c8092798 \
+    --hash=sha256:25a1f88b471b3bc911d18b935ecb7115dff3a192b6fef46f0bfaf71ff4f12418 \
+    --hash=sha256:25c8d773a62ce0451b020c7b29a35cfbc05de8b291163a7a0f3b7904f27253e6 \
+    --hash=sha256:2a60ad3e2553a74168d275a0ef35e8c0a965448ffbc3b300ab3a5bb9956c2162 \
+    --hash=sha256:2a66df3d4992bc1d725b9aa803e8c5a66c010c65c741ad901e260ece77f58d2f \
+    --hash=sha256:2ccc28197af5313706511fab3a8b66dcd6da067a1331372c82ea1cb74285e036 \
+    --hash=sha256:2e900bad2a8456d00a113cad8c13343f3b1f327534e3589acc2219729237a2e8 \
+    --hash=sha256:2ee7606193fb267be4b2e3b32714f2d58cad27217638db98a60f9efb5efeccc2 \
+    --hash=sha256:33ac8f098df0585c0b53009f039dfd913b38c1d2edafed0cedcc0c32a05aa110 \
+    --hash=sha256:3444cdba6628accf384e349014084b1cacd866fbb88433cd9d279d90a54e0b23 \
+    --hash=sha256:363ea8cd3c5cb6679f1c2f5f1f9669587361c062e4899fce56758efa928728f8 \
+    --hash=sha256:375a12d7556d462dc64d70475a9ee5982465fbb3d2b364f16b86ba9135793638 \
+    --hash=sha256:388f3217649d6d59292b722d940d4d2e1e6a7003259eb835724092a1cca0203a \
+    --hash=sha256:3947483a381259c06921612550867b37d22e1df6d6d7e8361264b6d037595f44 \
+    --hash=sha256:39e104da444a34830751715f45ef9fc537475ba21b7f1f5b0f4d71a3b60d7fe2 \
+    --hash=sha256:3c997f8c44ec9b9b0bcbf2d422cc00a1d9b9c681f56efa6ca149a941e5560da2 \
+    --hash=sha256:3dfafb44f7bb35c0c06eda6b2ab4bfd58f02729e7c4045e179f9a861b07c9850 \
+    --hash=sha256:3ebbcf2a07621f29638799828b8d8668c421bfb94c6cb04269130d8de4fb7136 \
+    --hash=sha256:3f88a4095e913f98988f5b338c1d4d5d07dbb0b6bad19892fd447484e483ba6b \
+    --hash=sha256:439e76255daa0f8151d3cb325f6dd4a3e93043e6403e6491813bcaaaa8733887 \
+    --hash=sha256:4569158070180c3855e9c0791c56be3ceeb192defa2cdf6a3f39e54319e56b89 \
+    --hash=sha256:466c219deee4536fbc83c08d09115249db301550625c7fef1c5563a584c9bc87 \
+    --hash=sha256:4a9d9b4d0a9b38d1c391bb4ad24aa65f306c6f01b512e10a8a34a2dc5675d348 \
+    --hash=sha256:4c7dde9e533c0a49d802b4f3f218fa9ad0a1ce21f2c2eb80d5216565202acab4 \
+    --hash=sha256:53d1bd3f979ed529f0805dd35ddaca330f80a9a6d90bc0121d2ff398f8ed8861 \
+    --hash=sha256:55346705687dbd7ef0d77883ab4f6fabc48232f587925bdaf95219bae072491e \
+    --hash=sha256:56295eb1e5f3aecd516d91b00cfd8bf3a13991de5a479df9e27dd569ea23959c \
+    --hash=sha256:56bb5c98f058a41bb58eead194b4db8c05b088c93d94d5161728515bd52b052b \
+    --hash=sha256:5a5b3bb545ead161be780ee85a2b54fdf7092815995661947812dde94a40f6fb \
+    --hash=sha256:5f2564ec89058ee7c7989a7b719115bdfe2a2fb8e7a4543b8d1c0cc4cf6478c1 \
+    --hash=sha256:608cce1da6f2672a56b24a015b42db4ac612ee709f3d29f27a00c943d9e851de \
+    --hash=sha256:63f13bf09cc3336eb04a837490b8f332e0db41da66995c9fd1ba04552e516354 \
+    --hash=sha256:662dd62358bdeaca0aee5761de8727cfd6861432e3bb828dc2a693aa0471a563 \
+    --hash=sha256:676135dcf3262c9c5081cc8f19ad55c8a64e3f7282a21266d05544450bffc3a5 \
+    --hash=sha256:67aeb72e0f482709991aa91345a831d0b707d16b0257e8ef88a2ad246a7280bf \
+    --hash=sha256:67b69535c870670c9f9b14a75d28baa32221d06f6b6fa6f77a0a13c5a7b0a5b9 \
+    --hash=sha256:682a7c79a2fbf40f5dbb1eb6bfe2cd865376deeac65acf9beb607505dced9e12 \
+    --hash=sha256:6994984550eaf25dd7fc7bd1b700ff45c894149341725bb4edc67f0ffa94efa4 \
+    --hash=sha256:69d3a98eebae99a420d4b28756c8ce6ea5a29291baf2dc9ff9414b42676f61d5 \
+    --hash=sha256:6e2e54267980349b723cff366d1e29b138b9a60fa376664a157a342689553f71 \
+    --hash=sha256:73e4b40ea0eda421b115248d7e79b59214411109a5bc47d0d48e4c73e3b8fcf9 \
+    --hash=sha256:74acd6e291f885678631b7ebc85d2d4aec458dd849b8c841b57ef04047833bed \
+    --hash=sha256:7665f04d0c7f26ff8bb534e1c65068409bf4687aa2534faf7104d7182debb336 \
+    --hash=sha256:7735e82e3498c27bcb2d17cb65d62c14f1100b71723b68362872bca7d0913d90 \
+    --hash=sha256:77a86c261679ea5f3896ec060be9dc8e365788248cc1e049632a1be682442063 \
+    --hash=sha256:7cf18abf9764746b9c8704774d8b06714bcb0a63641518a3a89c7f85cc02c2ad \
+    --hash=sha256:83928404adf8fb3d26793665633ea79b7361efa0287dfbd372a7e74311d51ee6 \
+    --hash=sha256:8e40876731f99b6f3c897b66b803c9e1c07a989b366c6b5b475fafd1f7ba3fb8 \
+    --hash=sha256:8f188cfcc64fb1266f4684206c9de0e80f54622c3f22a910cbd200478aeae61e \
+    --hash=sha256:91997d9cb4a325b60d4e3f20967f8eb08dfcb32b22554d5ef78e6fd1dda743a2 \
+    --hash=sha256:91ee8fc02ca52e24bcb77b234f22afc03288e1dafbb1f88fe24db308910c4ac7 \
+    --hash=sha256:92fe151145a990c22cbccf9ae15cae8ae9eddabfc949a219c9f667877e40853d \
+    --hash=sha256:945db8ee295d3af9dbdbb698cce9bbc5c59b5c3fe328bbc4387f59a8a35f998d \
+    --hash=sha256:9517d5e9e0731957468c29dbfd0f976736a0e55afaea843726e887f36fe017df \
+    --hash=sha256:952e0d9d07609d9c5be361f33b0d6d650cd2bae393aabb11d9b719364521984b \
+    --hash=sha256:97a58a28bcf63284e8b4d7b460cbee1edaab24634e82059c7b8c09e65284f178 \
+    --hash=sha256:97e48e8875e6c13909c800fa344cd54cc4b2b0db1d5f911f840458a500fde2c2 \
+    --hash=sha256:9e0f07b42d2a50c7dd2d8675d50f7343d998c64008f1da5fef888396b7f84630 \
+    --hash=sha256:a3dc1a4b165283bd865e8f8cb5f0c64c05001e0718ed06250d8cac9bec115b48 \
+    --hash=sha256:a3ebe9a75be7ab0b7da2464a77bb27febcb4fab46a34f9288f39d74833db7f61 \
+    --hash=sha256:a64e32f8bd94c105cc27f42d3b658902b5bcc947ece3c8fe7bc1b05982f60e89 \
+    --hash=sha256:a6ed8db0a556343d566a5c124ee483ae113acc9a557a807d439bcecc44e7dfbb \
+    --hash=sha256:ad9c9b99b05f163109466638bd30ada1722abb01bbb85c739c50b6dc11f92dc3 \
+    --hash=sha256:b33d7a286c0dc1a15f5fc864cc48ae92a846df287ceac2dd499926c3801054a6 \
+    --hash=sha256:bc092ba439d91df90aea38168e11f75c655880c12782facf5cf9c00f3d42b562 \
+    --hash=sha256:c436130cc779806bdf5d5fae0d848713105472b8566b75ff70048c47d3961c5b \
+    --hash=sha256:c5869b8fd70b81835a6f187c5fdbe67917a04d7e52b6e7cc4e5fe39d55c39d58 \
+    --hash=sha256:c5ecca8f9bab618340c8e848d340baf68bcd8ad90a8ecd7a4524a81c1764b3db \
+    --hash=sha256:cfac69017ef97db2438efb854edf24f5a29fd09a536ff3a992b75990720cdc99 \
+    --hash=sha256:d2f0d0f976985f85dfb5f3d685697ef769faa6b71993b46b295cdbbd6be8cc37 \
+    --hash=sha256:d5bed7f9805cc29c780f3aee05de3262ee7ce1f47083cfe9f77471e9d6777e83 \
+    --hash=sha256:d6a21ef516d36909931a2967621eecb256018aeb11fc48656e3257e73e2e247a \
+    --hash=sha256:d9b6ddac6408194e934002a69bcaadbc88c10b5f38fb9307779d1c629181815d \
+    --hash=sha256:db47514ffdbd91ccdc7e6f8407aac4ee94cc871b15b577c1c324236b013ddd04 \
+    --hash=sha256:df81779732feb9d01e5d513fad0122efb3d53bbc75f61b2a4f29a020bc985e70 \
+    --hash=sha256:e4a91d44379f45f5e540971d41e4626dacd7f01004826a18cb048e7da7e96544 \
+    --hash=sha256:e63e3e1e0271f374ed489ff5ee73d4b6e7c60710e1f76af5f0e1a6117cd26394 \
+    --hash=sha256:e70fac33e8b4ac63dfc4c956fd7d85a0b1139adcfc0d964ce288b7c527537fea \
+    --hash=sha256:ecddc221a077a8132cf7c747d5352a15ed763b674c0448d811f408bf803d9ad7 \
+    --hash=sha256:f45eec587dafd4b2d41ac189c2156461ebd0c1082d2fe7013571598abb8505d1 \
+    --hash=sha256:f52a68c21363c45297aca15561812d542f8fc683c85201df0bebe209e349f793 \
+    --hash=sha256:f571aea50ba5623c308aa146eb650eebf7dbe0fd8c5d946e28343cb3b5aad577 \
+    --hash=sha256:f60f0ac7005b9f5a6091009b09a419ace1610e163fa5deaba5ce3484341840e7 \
+    --hash=sha256:f6475a1b2ecb310c98c28d271a30df74f9dd436ee46d09236a6b750a7599ce57 \
+    --hash=sha256:f6d5749fdd33d90e34c2efb174c7e236829147a2713334d708746e94c4bde40d \
+    --hash=sha256:f902804113e032e2cdf8c71015651c97af6418363bea8d78dc0911d56c335032 \
+    --hash=sha256:fa1076244f54bb76e65e22cb6910365779d5c3d71d1f18b275f1dfc7b0d71b4d \
+    --hash=sha256:fc2db02409338bf36590aa985a461b2c96fce91f8e7e0f14c50c5fcc4f229016 \
+    --hash=sha256:ffcad6c564fe6b9b8916c1aefbb37a362deebf9394bd2974e9d84232e3e08504
+    # via yarl
+proto-plus==1.25.0 \
+    --hash=sha256:c91fc4a65074ade8e458e95ef8bac34d4008daa7cce4a12d6707066fca648961 \
+    --hash=sha256:fbb17f57f7bd05a68b7707e745e26528b0b3c34e378db91eef93912c54982d91
+    # via google-api-core
 protobuf==5.28.2 \
     --hash=sha256:2c69461a7fcc8e24be697624c09a839976d82ae75062b11a0972e41fd2cd9132 \
     --hash=sha256:35cfcb15f213449af7ff6198d6eb5f739c37d7e4f1c09b5d0641babf2cc0c68f \
@@ -523,11 +1197,33 @@ protobuf==5.28.2 \
     --hash=sha256:c0ea0123dac3399a2eeb1a1443d82b7afc9ff40241433296769f7da42d142ec3 \
     --hash=sha256:ca53faf29896c526863366a52a8f4d88e69cd04ec9571ed6082fa117fac3ab36 \
     --hash=sha256:eeea10f3dc0ac7e6b4933d32db20662902b4ab81bf28df12218aa389e9c2102d
-    # via orbax-checkpoint
+    # via
+    #   google-api-core
+    #   googleapis-common-protos
+    #   orbax-checkpoint
+    #   proto-plus
+pyasn1==0.6.1 \
+    --hash=sha256:0d632f46f2ba09143da3a8afe9e33fb6f92fa2320ab7e886e2d0f7672af84629 \
+    --hash=sha256:6f580d2bdd84365380830acf45550f2511469f673cb4a5ae3857a3170128b034
+    # via
+    #   pyasn1-modules
+    #   rsa
+pyasn1-modules==0.4.1 \
+    --hash=sha256:49bfa96b45a292b711e986f222502c1c9a5e1f4e568fc30e2574a6c7d07838fd \
+    --hash=sha256:c28e2dbf9c06ad61c71a075c7e0f9fd0f1b0bb2d2ad4377f240d33ac2ab60a7c
+    # via google-auth
 pygments==2.18.0 \
     --hash=sha256:786ff802f32e91311bff3889f6e9a86e81505fe99f2735bb6d60ae0c5004f199 \
     --hash=sha256:b8e6aca0523f3ab76fee51799c488e38782ac06eafcf95e7ba832985c8e7b13a
     # via rich
+python-dateutil==2.9.0.post0 \
+    --hash=sha256:37dd54208da7e1cd875388217d5e00ebd4179249f90fb72437e91a35459a0ad3 \
+    --hash=sha256:a8b2bc7bffae282281c8140a97d3aa9c14da0b136dfe83f850eea9a5f7470427
+    # via pandas
+pytz==2024.2 \
+    --hash=sha256:2aa355083c50a0f93fa581709deac0c9ad65cca8a9e9beac660adcbd493c798a \
+    --hash=sha256:31c7c1817eb7fae7ca4b8c7ee50c72f93aa2dd863de768e1ef4245d426aa0725
+    # via pandas
 pyyaml==6.0.2 \
     --hash=sha256:01179a4a8559ab5de078078f37e5c1a30d76bb88519906844fd7bdea1b7729ff \
     --hash=sha256:0833f8694549e586547b576dcfaba4a6b55b9e96098b36cdc7ebefe667dfed48 \
@@ -589,11 +1285,52 @@ pyyaml==6.0.2 \
 requests==2.32.3 ; sys_platform == "linux" \
     --hash=sha256:55365417734eb18255590a9ff9eb97e9e1da868d4ccd6402399eaf68af20a760 \
     --hash=sha256:70761cfe03c773ceb22aa2f671b4757976145175cdfca038c02654d061d6dcc6
-    # via -r builddeps/test-requirements.txt
+    # via
+    #   -r builddeps/test-requirements.txt
+    #   gcsfs
+    #   google-api-core
+    #   google-cloud-storage
+    #   requests-oauthlib
+requests-oauthlib==2.0.0 \
+    --hash=sha256:7dd8a5c40426b779b0868c404bdef9768deccf22749cde15852df527e6269b36 \
+    --hash=sha256:b3dffaebd884d8cd778494369603a9e7b58d29111bf6b41bdc2dcd87203af4e9
+    # via google-auth-oauthlib
 rich==13.9.2 \
     --hash=sha256:51a2c62057461aaf7152b4d611168f93a9fc73068f8ded2790f29fe2b5366d0c \
     --hash=sha256:8c82a3d3f8dcfe9e734771313e606b39d8247bb6b826e196f4914b333b743cf1
     # via flax
+rsa==4.9 \
+    --hash=sha256:90260d9058e514786967344d0ef75fa8727eed8a7d2e43ce9f4bcf1b536174f7 \
+    --hash=sha256:e38464a49c6c85d7f1351b0126661487a7e0a14a50f1675ec50eb34d4f20ef21
+    # via google-auth
+scikit-learn==1.5.2 \
+    --hash=sha256:03b6158efa3faaf1feea3faa884c840ebd61b6484167c711548fce208ea09445 \
+    --hash=sha256:178ddd0a5cb0044464fc1bfc4cca5b1833bfc7bb022d70b05db8530da4bb3dd3 \
+    --hash=sha256:1ff45e26928d3b4eb767a8f14a9a6efbf1cbff7c05d1fb0f95f211a89fd4f5de \
+    --hash=sha256:299406827fb9a4f862626d0fe6c122f5f87f8910b86fe5daa4c32dcd742139b6 \
+    --hash=sha256:2d4cad1119c77930b235579ad0dc25e65c917e756fe80cab96aa3b9428bd3fb0 \
+    --hash=sha256:394397841449853c2290a32050382edaec3da89e35b3e03d6cc966aebc6a8ae6 \
+    --hash=sha256:3a686885a4b3818d9e62904d91b57fa757fc2bed3e465c8b177be652f4dd37c8 \
+    --hash=sha256:3b923d119d65b7bd555c73be5423bf06c0105678ce7e1f558cb4b40b0a5502b1 \
+    --hash=sha256:3bed4909ba187aca80580fe2ef370d9180dcf18e621a27c4cf2ef10d279a7efe \
+    --hash=sha256:52788f48b5d8bca5c0736c175fa6bdaab2ef00a8f536cda698db61bd89c551c1 \
+    --hash=sha256:57cc1786cfd6bd118220a92ede80270132aa353647684efa385a74244a41e3b1 \
+    --hash=sha256:643964678f4b5fbdc95cbf8aec638acc7aa70f5f79ee2cdad1eec3df4ba6ead8 \
+    --hash=sha256:6c16d84a0d45e4894832b3c4d0bf73050939e21b99b01b6fd59cbb0cf39163b6 \
+    --hash=sha256:757c7d514ddb00ae249832fe87100d9c73c6ea91423802872d9e74970a0e40b9 \
+    --hash=sha256:8c412ccc2ad9bf3755915e3908e677b367ebc8d010acbb3f182814524f2e5540 \
+    --hash=sha256:b0768ad641981f5d3a198430a1d31c3e044ed2e8a6f22166b4d546a5116d7908 \
+    --hash=sha256:b4237ed7b3fdd0a4882792e68ef2545d5baa50aca3bb45aa7df468138ad8f94d \
+    --hash=sha256:b7b0f9a0b1040830d38c39b91b3a44e1b643f4b36e36567b80b7c6bd2202a27f \
+    --hash=sha256:c15b1ca23d7c5f33cc2cb0a0d6aaacf893792271cddff0edbd6a40e8319bc113 \
+    --hash=sha256:ca64b3089a6d9b9363cd3546f8978229dcbb737aceb2c12144ee3f70f95684b7 \
+    --hash=sha256:e9a702e2de732bbb20d3bad29ebd77fc05a6b427dc49964300340e4c9328b3f5 \
+    --hash=sha256:f60021ec1574e56632be2a36b946f8143bf4e5e6af4a06d85281adc22938e0dd \
+    --hash=sha256:f7284ade780084d94505632241bf78c44ab3b6f1e8ccab3d2af58e0e950f9c12 \
+    --hash=sha256:f763897fe92d0e903aa4847b0aec0e68cadfff77e8a0687cabd946c89d17e675 \
+    --hash=sha256:f8b0ccd4a902836493e026c03256e8b206656f91fbcc4fde28c57a5b752561f1 \
+    --hash=sha256:f932a02c3f4956dfb981391ab24bda1dbd90fe3d628e4b42caef3e041c67707a
+    # via dinosaur-dycore
 scipy==1.14.1 \
     --hash=sha256:0c2f95de3b04e26f5f3ad5bb05e74ba7f68b837133a4492414b3afd79dfe540e \
     --hash=sha256:1729560c906963fc8389f6aac023739ff3983e727b1a4d87696b7bf108316a79 \
@@ -629,12 +1366,17 @@ scipy==1.14.1 \
     --hash=sha256:edaf02b82cd7639db00dbff629995ef185c8df4c3ffa71a5562a595765a06ce1 \
     --hash=sha256:fef8c87f8abfb884dac04e97824b61299880c43f4ce675dd2cbeadd3c9b466d2
     # via
+    #   dinosaur-dycore
     #   jax
     #   jaxlib
+    #   scikit-learn
 six==1.16.0 \
     --hash=sha256:1e61c37477a1626458e36f7b1d82aa5c9b094fa4802892072e49de9c60c4c926 \
     --hash=sha256:8abb2f1d86890a2dfb989f9a77cfcfd3e47c2a354b01111771326f8aa26e0254
-    # via ml-collections
+    # via
+    #   ml-collections
+    #   python-dateutil
+    #   tensorflow-probability
 sympy==1.13.3 \
     --hash=sha256:54612cf55a62755ee71824ce692986f23c88ffa77207b30c1368eda4a7060f73 \
     --hash=sha256:b27fd2c6530e0ab39e275fc9b683895367e51d5da91baa8d3d64db2565fec4d9
@@ -643,6 +1385,9 @@ tabulate==0.9.0 \
     --hash=sha256:0095b12bf5966de529c0feb1fa08671671b3368eec77d7ef7ab114be2c068b3c \
     --hash=sha256:024ca478df22e9340661486f85298cff5f6dcdba14f3813e8830015b9ed1948f
     # via dm-haiku
+tensorflow-probability==0.24.0 \
+    --hash=sha256:8c1774683e38359dbcaf3697e79b7e6a4e69b9c7b3679e78ee18f43e59e5759b
+    # via neuralgcm
 tensorstore==0.1.66 \
     --hash=sha256:0261b1b7121186ac4898c9b4e29289c46754f9d7d73e906620e9a21386fe564e \
     --hash=sha256:06f7dba630d75f65201b1c021fcd89e5f8538506b444b5dd712f35c11b812fc8 \
@@ -668,10 +1413,19 @@ tensorstore==0.1.66 \
     # via
     #   flax
     #   orbax-checkpoint
+    #   xarray-tensorstore
+threadpoolctl==3.5.0 \
+    --hash=sha256:082433502dd922bf738de0d8bcc4fdcbf0979ff44c42bd40f5af8a282f6fa107 \
+    --hash=sha256:56c1e26c150397e58c4926da8eeee87533b1e32bef131bd4bf6a2f45f3185467
+    # via scikit-learn
 toolz==1.0.0 \
     --hash=sha256:292c8f1c4e7516bf9086f8850935c799a874039c8bcf959d47b600e4c44a6236 \
     --hash=sha256:2c86e3d9a04798ac556793bced838816296a2f085017664e4995cb40a1047a02
     # via chex
+tree-math==0.2.1 \
+    --hash=sha256:77613c2360e067b16c0738b32e7d64aa5e1167e609ca1a32b6fa05dcb0d330bb \
+    --hash=sha256:80b01a5d29c1339db0cd8b465f9fad4f872adac117ed599b14f093a6c52821fe
+    # via dinosaur-dycore
 typing-extensions==4.12.2 \
     --hash=sha256:04e5ca0351e0f3f85c6853954072df659d0d13fac324d0072316b67d7794700d \
     --hash=sha256:1a7ead55c7e559dd4dee8856e3a88b41225abfe1ce8df57b7c13915fe121ffb8
@@ -679,12 +1433,119 @@ typing-extensions==4.12.2 \
     #   chex
     #   etils
     #   flax
+    #   flexcache
+    #   flexparser
+    #   multidict
     #   orbax-checkpoint
+    #   pint
     #   rich
+tzdata==2024.2 \
+    --hash=sha256:7d85cc416e9382e69095b7bdf4afd9e3880418a2413feec7069d533d6b4e31cc \
+    --hash=sha256:a48093786cdcde33cad18c2555e8532f34422074448fbc874186f0abd79565cd
+    # via pandas
 urllib3==2.2.3 \
     --hash=sha256:ca899ca043dcb1bafa3e262d73aa25c465bfb49e0bd9dd5d59f1d0acba2f8fac \
     --hash=sha256:e7d814a81dad81e6caf2ec9fdedb284ecc9c73076b62654547cc64ccdcae26e9
     # via requests
+xarray==2024.10.0 \
+    --hash=sha256:ae1d38cb44a0324dfb61e492394158ae22389bf7de9f3c174309c17376df63a0 \
+    --hash=sha256:e369e2bac430e418c2448e5b96f07da4635f98c1319aa23cfeb3fbcb9a01d2e0
+    # via
+    #   dinosaur-dycore
+    #   neuralgcm
+    #   xarray-tensorstore
+xarray-tensorstore==0.1.4 \
+    --hash=sha256:1dacec6376912373dbf9b239d96184090f400a49cd66017f5bf0c824793c9552 \
+    --hash=sha256:de3558c08891c4a7524cf579cb271d6cc005962efd1b8601dd679b58b13a68c8
+    # via dinosaur-dycore
+yarl==1.17.1 \
+    --hash=sha256:06157fb3c58f2736a5e47c8fcbe1afc8b5de6fb28b14d25574af9e62150fcaac \
+    --hash=sha256:067a63fcfda82da6b198fa73079b1ca40b7c9b7994995b6ee38acda728b64d47 \
+    --hash=sha256:0b1794853124e2f663f0ea54efb0340b457f08d40a1cef78edfa086576179c91 \
+    --hash=sha256:0bdff5e0995522706c53078f531fb586f56de9c4c81c243865dd5c66c132c3b5 \
+    --hash=sha256:117ed8b3732528a1e41af3aa6d4e08483c2f0f2e3d3d7dca7cf538b3516d93df \
+    --hash=sha256:14bc88baa44e1f84164a392827b5defb4fa8e56b93fecac3d15315e7c8e5d8b3 \
+    --hash=sha256:1654ec814b18be1af2c857aa9000de7a601400bd4c9ca24629b18486c2e35463 \
+    --hash=sha256:16bca6678a83657dd48df84b51bd56a6c6bd401853aef6d09dc2506a78484c7b \
+    --hash=sha256:1a3b91c44efa29e6c8ef8a9a2b583347998e2ba52c5d8280dbd5919c02dfc3b5 \
+    --hash=sha256:1a52a1ffdd824fb1835272e125385c32fd8b17fbdefeedcb4d543cc23b332d74 \
+    --hash=sha256:1ce36ded585f45b1e9bb36d0ae94765c6608b43bd2e7f5f88079f7a85c61a4d3 \
+    --hash=sha256:299f11b44d8d3a588234adbe01112126010bd96d9139c3ba7b3badd9829261c3 \
+    --hash=sha256:2b24ec55fad43e476905eceaf14f41f6478780b870eda5d08b4d6de9a60b65b4 \
+    --hash=sha256:2d374d70fdc36f5863b84e54775452f68639bc862918602d028f89310a034ab0 \
+    --hash=sha256:2d9f0606baaec5dd54cb99667fcf85183a7477f3766fbddbe3f385e7fc253299 \
+    --hash=sha256:2e7ba4c9377e48fb7b20dedbd473cbcbc13e72e1826917c185157a137dac9df2 \
+    --hash=sha256:2f0a6423295a0d282d00e8701fe763eeefba8037e984ad5de44aa349002562ac \
+    --hash=sha256:327828786da2006085a4d1feb2594de6f6d26f8af48b81eb1ae950c788d97f61 \
+    --hash=sha256:380e6c38ef692b8fd5a0f6d1fa8774d81ebc08cfbd624b1bca62a4d4af2f9931 \
+    --hash=sha256:3b74ff4767d3ef47ffe0cd1d89379dc4d828d4873e5528976ced3b44fe5b0a21 \
+    --hash=sha256:3e844be8d536afa129366d9af76ed7cb8dfefec99f5f1c9e4f8ae542279a6dc3 \
+    --hash=sha256:459e81c2fb920b5f5df744262d1498ec2c8081acdcfe18181da44c50f51312f7 \
+    --hash=sha256:46ddf6e0b975cd680eb83318aa1d321cb2bf8d288d50f1754526230fcf59ba96 \
+    --hash=sha256:482c122b72e3c5ec98f11457aeb436ae4aecca75de19b3d1de7cf88bc40db82f \
+    --hash=sha256:561c87fea99545ef7d692403c110b2f99dced6dff93056d6e04384ad3bc46243 \
+    --hash=sha256:578d00c9b7fccfa1745a44f4eddfdc99d723d157dad26764538fbdda37209857 \
+    --hash=sha256:58c8e9620eb82a189c6c40cb6b59b4e35b2ee68b1f2afa6597732a2b467d7e8f \
+    --hash=sha256:5b29beab10211a746f9846baa39275e80034e065460d99eb51e45c9a9495bcca \
+    --hash=sha256:5d1d42556b063d579cae59e37a38c61f4402b47d70c29f0ef15cee1acaa64488 \
+    --hash=sha256:5f236cb5999ccd23a0ab1bd219cfe0ee3e1c1b65aaf6dd3320e972f7ec3a39da \
+    --hash=sha256:62a91aefff3d11bf60e5956d340eb507a983a7ec802b19072bb989ce120cd948 \
+    --hash=sha256:64cc6e97f14cf8a275d79c5002281f3040c12e2e4220623b5759ea7f9868d6a5 \
+    --hash=sha256:6f4c9156c4d1eb490fe374fb294deeb7bc7eaccda50e23775b2354b6a6739934 \
+    --hash=sha256:7294e38f9aa2e9f05f765b28ffdc5d81378508ce6dadbe93f6d464a8c9594473 \
+    --hash=sha256:7615058aabad54416ddac99ade09a5510cf77039a3b903e94e8922f25ed203d7 \
+    --hash=sha256:7e48cdb8226644e2fbd0bdb0a0f87906a3db07087f4de77a1b1b1ccfd9e93685 \
+    --hash=sha256:7f63d176a81555984e91f2c84c2a574a61cab7111cc907e176f0f01538e9ff6e \
+    --hash=sha256:7f6595c852ca544aaeeb32d357e62c9c780eac69dcd34e40cae7b55bc4fb1147 \
+    --hash=sha256:7fac95714b09da9278a0b52e492466f773cfe37651cf467a83a1b659be24bf71 \
+    --hash=sha256:81713b70bea5c1386dc2f32a8f0dab4148a2928c7495c808c541ee0aae614d67 \
+    --hash=sha256:846dd2e1243407133d3195d2d7e4ceefcaa5f5bf7278f0a9bda00967e6326b04 \
+    --hash=sha256:84c063af19ef5130084db70ada40ce63a84f6c1ef4d3dbc34e5e8c4febb20822 \
+    --hash=sha256:881764d610e3269964fc4bb3c19bb6fce55422828e152b885609ec176b41cf11 \
+    --hash=sha256:8994b29c462de9a8fce2d591028b986dbbe1b32f3ad600b2d3e1c482c93abad6 \
+    --hash=sha256:8c79e9d7e3d8a32d4824250a9c6401194fb4c2ad9a0cec8f6a96e09a582c2cc0 \
+    --hash=sha256:8ee427208c675f1b6e344a1f89376a9613fc30b52646a04ac0c1f6587c7e46ec \
+    --hash=sha256:949681f68e0e3c25377462be4b658500e85ca24323d9619fdc41f68d46a1ffda \
+    --hash=sha256:9e275792097c9f7e80741c36de3b61917aebecc08a67ae62899b074566ff8556 \
+    --hash=sha256:9fb815155aac6bfa8d86184079652c9715c812d506b22cfa369196ef4e99d1b4 \
+    --hash=sha256:a2a64e62c7a0edd07c1c917b0586655f3362d2c2d37d474db1a509efb96fea1c \
+    --hash=sha256:a7ac5b4984c468ce4f4a553df281450df0a34aefae02e58d77a0847be8d1e11f \
+    --hash=sha256:aa46dce75078fceaf7cecac5817422febb4355fbdda440db55206e3bd288cfb8 \
+    --hash=sha256:ae3476e934b9d714aa8000d2e4c01eb2590eee10b9d8cd03e7983ad65dfbfcba \
+    --hash=sha256:b0341e6d9a0c0e3cdc65857ef518bb05b410dbd70d749a0d33ac0f39e81a4258 \
+    --hash=sha256:b40d1bf6e6f74f7c0a567a9e5e778bbd4699d1d3d2c0fe46f4b717eef9e96b95 \
+    --hash=sha256:b5c4804e4039f487e942c13381e6c27b4b4e66066d94ef1fae3f6ba8b953f383 \
+    --hash=sha256:b5d6a6c9602fd4598fa07e0389e19fe199ae96449008d8304bf5d47cb745462e \
+    --hash=sha256:b5f1ac7359e17efe0b6e5fec21de34145caef22b260e978336f325d5c84e6938 \
+    --hash=sha256:c0167540094838ee9093ef6cc2c69d0074bbf84a432b4995835e8e5a0d984374 \
+    --hash=sha256:c180ac742a083e109c1a18151f4dd8675f32679985a1c750d2ff806796165b55 \
+    --hash=sha256:c73df5b6e8fabe2ddb74876fb82d9dd44cbace0ca12e8861ce9155ad3c886139 \
+    --hash=sha256:c7e177c619342e407415d4f35dec63d2d134d951e24b5166afcdfd1362828e17 \
+    --hash=sha256:cbad927ea8ed814622305d842c93412cb47bd39a496ed0f96bfd42b922b4a217 \
+    --hash=sha256:cc353841428d56b683a123a813e6a686e07026d6b1c5757970a877195f880c2d \
+    --hash=sha256:cc7c92c1baa629cb03ecb0c3d12564f172218fb1739f54bf5f3881844daadc6d \
+    --hash=sha256:cc7d768260f4ba4ea01741c1b5fe3d3a6c70eb91c87f4c8761bbcce5181beafe \
+    --hash=sha256:d0eea830b591dbc68e030c86a9569826145df485b2b4554874b07fea1275a199 \
+    --hash=sha256:d216e5d9b8749563c7f2c6f7a0831057ec844c68b4c11cb10fc62d4fd373c26d \
+    --hash=sha256:d401f07261dc5aa36c2e4efc308548f6ae943bfff20fcadb0a07517a26b196d8 \
+    --hash=sha256:d6324274b4e0e2fa1b3eccb25997b1c9ed134ff61d296448ab8269f5ac068c4c \
+    --hash=sha256:d8a8b74d843c2638f3864a17d97a4acda58e40d3e44b6303b8cc3d3c44ae2d29 \
+    --hash=sha256:d9b6b28a57feb51605d6ae5e61a9044a31742db557a3b851a74c13bc61de5172 \
+    --hash=sha256:de599af166970d6a61accde358ec9ded821234cbbc8c6413acfec06056b8e860 \
+    --hash=sha256:e594b22688d5747b06e957f1ef822060cb5cb35b493066e33ceac0cf882188b7 \
+    --hash=sha256:e5b078134f48552c4d9527db2f7da0b5359abd49393cdf9794017baec7506170 \
+    --hash=sha256:eb6dce402734575e1a8cc0bb1509afca508a400a57ce13d306ea2c663bad1138 \
+    --hash=sha256:f1790a4b1e8e8e028c391175433b9c8122c39b46e1663228158e61e6f915bf06 \
+    --hash=sha256:f5efe0661b9fcd6246f27957f6ae1c0eb29bc60552820f01e970b4996e016004 \
+    --hash=sha256:f9cbfbc5faca235fbdf531b93aa0f9f005ec7d267d9d738761a4d42b744ea159 \
+    --hash=sha256:fbea1751729afe607d84acfd01efd95e3b31db148a181a441984ce9b3d3469da \
+    --hash=sha256:fca4b4307ebe9c3ec77a084da3a9d1999d164693d16492ca2b64594340999988 \
+    --hash=sha256:ff5c6771c7e3511a06555afa317879b7db8d640137ba55d6ab0d0c50425cab75
+    # via aiohttp
+zarr==2.18.3 \
+    --hash=sha256:2580d8cb6dd84621771a10d31c4d777dca8a27706a1a89b29f42d2d37e2df5ce \
+    --hash=sha256:b1f7dfd2496f436745cdd4c7bcf8d3b4bc1dceef5fdd0d589c87130d842496dd
+    # via xarray-tensorstore
 zipp==3.20.2 \
     --hash=sha256:a817ac80d6cf4b23bf7f2828b7cabf326f15a001bea8b1f9b49631780ba28350 \
     --hash=sha256:bc9eb26f4506fda01b81bcde0ca78103b6e62f991b381fec825435c836edbc29
diff --git a/builddeps/requirements_lock_3_11.txt b/builddeps/requirements_lock_3_11.txt
index 7e14f174..d74a4a59 100644
--- a/builddeps/requirements_lock_3_11.txt
+++ b/builddeps/requirements_lock_3_11.txt
@@ -16,10 +16,121 @@ absl-py==2.1.0 \
     #   ml-collections
     #   optax
     #   orbax-checkpoint
+    #   tensorflow-probability
+aiohappyeyeballs==2.4.3 \
+    --hash=sha256:75cf88a15106a5002a8eb1dab212525c00d1f4c0fa96e551c9fbe6f09a621586 \
+    --hash=sha256:8a7a83727b2756f394ab2895ea0765a0a8c475e3c71e98d43d76f22b4b435572
+    # via aiohttp
+aiohttp==3.10.10 \
+    --hash=sha256:007ec22fbc573e5eb2fb7dec4198ef8f6bf2fe4ce20020798b2eb5d0abda6138 \
+    --hash=sha256:00819de9e45d42584bed046314c40ea7e9aea95411b38971082cad449392b08c \
+    --hash=sha256:01948b1d570f83ee7bbf5a60ea2375a89dfb09fd419170e7f5af029510033d24 \
+    --hash=sha256:038f514fe39e235e9fef6717fbf944057bfa24f9b3db9ee551a7ecf584b5b480 \
+    --hash=sha256:03a42ac7895406220124c88911ebee31ba8b2d24c98507f4a8bf826b2937c7f2 \
+    --hash=sha256:05646ebe6b94cc93407b3bf34b9eb26c20722384d068eb7339de802154d61bc5 \
+    --hash=sha256:0631dd7c9f0822cc61c88586ca76d5b5ada26538097d0f1df510b082bad3411a \
+    --hash=sha256:0b00807e2605f16e1e198f33a53ce3c4523114059b0c09c337209ae55e3823a8 \
+    --hash=sha256:0e1b370d8007c4ae31ee6db7f9a2fe801a42b146cec80a86766e7ad5c4a259cf \
+    --hash=sha256:15ecd889a709b0080f02721255b3f80bb261c2293d3c748151274dfea93ac871 \
+    --hash=sha256:1b66ccafef7336a1e1f0e389901f60c1d920102315a56df85e49552308fc0486 \
+    --hash=sha256:1bbb122c557a16fafc10354b9d99ebf2f2808a660d78202f10ba9d50786384b9 \
+    --hash=sha256:1eb89d3d29adaf533588f209768a9c02e44e4baf832b08118749c5fad191781d \
+    --hash=sha256:258c5dd01afc10015866114e210fb7365f0d02d9d059c3c3415382ab633fcbcb \
+    --hash=sha256:2609e9ab08474702cc67b7702dbb8a80e392c54613ebe80db7e8dbdb79837c68 \
+    --hash=sha256:274cfa632350225ce3fdeb318c23b4a10ec25c0e2c880eff951a3842cf358ac1 \
+    --hash=sha256:28529e08fde6f12eba8677f5a8608500ed33c086f974de68cc65ab218713a59d \
+    --hash=sha256:2b606353da03edcc71130b52388d25f9a30a126e04caef1fd637e31683033abd \
+    --hash=sha256:30ca7c3b94708a9d7ae76ff281b2f47d8eaf2579cd05971b5dc681db8caac6e1 \
+    --hash=sha256:333cf6cf8e65f6a1e06e9eb3e643a0c515bb850d470902274239fea02033e9a8 \
+    --hash=sha256:3455522392fb15ff549d92fbf4b73b559d5e43dc522588f7eb3e54c3f38beee7 \
+    --hash=sha256:362f641f9071e5f3ee6f8e7d37d5ed0d95aae656adf4ef578313ee585b585959 \
+    --hash=sha256:3bcd391d083f636c06a68715e69467963d1f9600f85ef556ea82e9ef25f043f7 \
+    --hash=sha256:3dffb610a30d643983aeb185ce134f97f290f8935f0abccdd32c77bed9388b42 \
+    --hash=sha256:3fe407bf93533a6fa82dece0e74dbcaaf5d684e5a51862887f9eaebe6372cd79 \
+    --hash=sha256:413251f6fcf552a33c981c4709a6bba37b12710982fec8e558ae944bfb2abd38 \
+    --hash=sha256:438cd072f75bb6612f2aca29f8bd7cdf6e35e8f160bc312e49fbecab77c99e3a \
+    --hash=sha256:4470c73c12cd9109db8277287d11f9dd98f77fc54155fc71a7738a83ffcc8ea8 \
+    --hash=sha256:45c3b868724137f713a38376fef8120c166d1eadd50da1855c112fe97954aed8 \
+    --hash=sha256:486f7aabfa292719a2753c016cc3a8f8172965cabb3ea2e7f7436c7f5a22a151 \
+    --hash=sha256:4f05e9727ce409358baa615dbeb9b969db94324a79b5a5cea45d39bdb01d82e6 \
+    --hash=sha256:50aed5155f819873d23520919e16703fc8925e509abbb1a1491b0087d1cd969e \
+    --hash=sha256:50edbcad60d8f0e3eccc68da67f37268b5144ecc34d59f27a02f9611c1d4eec7 \
+    --hash=sha256:54ca74df1be3c7ca1cf7f4c971c79c2daf48d9aa65dea1a662ae18926f5bc8ce \
+    --hash=sha256:578a4b875af3e0daaf1ac6fa983d93e0bbfec3ead753b6d6f33d467100cdc67b \
+    --hash=sha256:597a079284b7ee65ee102bc3a6ea226a37d2b96d0418cc9047490f231dc09fe8 \
+    --hash=sha256:59bb3c54aa420521dc4ce3cc2c3fe2ad82adf7b09403fa1f48ae45c0cbde6628 \
+    --hash=sha256:5c6a5b8c7926ba5d8545c7dd22961a107526562da31a7a32fa2456baf040939f \
+    --hash=sha256:64f6c17757251e2b8d885d728b6433d9d970573586a78b78ba8929b0f41d045a \
+    --hash=sha256:679abe5d3858b33c2cf74faec299fda60ea9de62916e8b67e625d65bf069a3b7 \
+    --hash=sha256:741a46d58677d8c733175d7e5aa618d277cd9d880301a380fd296975a9cdd7bc \
+    --hash=sha256:7789050d9e5d0c309c706953e5e8876e38662d57d45f936902e176d19f1c58ab \
+    --hash=sha256:77abf6665ae54000b98b3c742bc6ea1d1fb31c394bcabf8b5d2c1ac3ebfe7f3b \
+    --hash=sha256:79019094f87c9fb44f8d769e41dbb664d6e8fcfd62f665ccce36762deaa0e911 \
+    --hash=sha256:7b06b7843929e41a94ea09eb1ce3927865387e3e23ebe108e0d0d09b08d25be9 \
+    --hash=sha256:7e338c0523d024fad378b376a79faff37fafb3c001872a618cde1d322400a572 \
+    --hash=sha256:7ea7ffc6d6d6f8a11e6f40091a1040995cdff02cfc9ba4c2f30a516cb2633554 \
+    --hash=sha256:8105fd8a890df77b76dd3054cddf01a879fc13e8af576805d667e0fa0224c35d \
+    --hash=sha256:84afcdea18eda514c25bc68b9af2a2b1adea7c08899175a51fe7c4fb6d551257 \
+    --hash=sha256:9294bbb581f92770e6ed5c19559e1e99255e4ca604a22c5c6397b2f9dd3ee42c \
+    --hash=sha256:93429602396f3383a797a2a70e5f1de5df8e35535d7806c9f91df06f297e109b \
+    --hash=sha256:9627cc1a10c8c409b5822a92d57a77f383b554463d1884008e051c32ab1b3742 \
+    --hash=sha256:998f3bd3cfc95e9424a6acd7840cbdd39e45bc09ef87533c006f94ac47296090 \
+    --hash=sha256:9c72109213eb9d3874f7ac8c0c5fa90e072d678e117d9061c06e30c85b4cf0e6 \
+    --hash=sha256:9fc1500fd2a952c5c8e3b29aaf7e3cc6e27e9cfc0a8819b3bce48cc1b849e4cc \
+    --hash=sha256:a3f00003de6eba42d6e94fabb4125600d6e484846dbf90ea8e48a800430cc142 \
+    --hash=sha256:a45d85cf20b5e0d0aa5a8dca27cce8eddef3292bc29d72dcad1641f4ed50aa16 \
+    --hash=sha256:a7d8d14fe962153fc681f6366bdec33d4356f98a3e3567782aac1b6e0e40109a \
+    --hash=sha256:a8fa23fe62c436ccf23ff930149c047f060c7126eae3ccea005f0483f27b2e28 \
+    --hash=sha256:aa6658732517ddabe22c9036479eabce6036655ba87a0224c612e1ae6af2087e \
+    --hash=sha256:aafc8ee9b742ce75044ae9a4d3e60e3d918d15a4c2e08a6c3c3e38fa59b92d94 \
+    --hash=sha256:ab5a5a0c7a7991d90446a198689c0535be89bbd6b410a1f9a66688f0880ec026 \
+    --hash=sha256:acd48d5b80ee80f9432a165c0ac8cbf9253eaddb6113269a5e18699b33958dbb \
+    --hash=sha256:ad7593bb24b2ab09e65e8a1d385606f0f47c65b5a2ae6c551db67d6653e78c28 \
+    --hash=sha256:baa42524a82f75303f714108fea528ccacf0386af429b69fff141ffef1c534f9 \
+    --hash=sha256:bdfcf6443637c148c4e1a20c48c566aa694fa5e288d34b20fcdc58507882fed3 \
+    --hash=sha256:be7443669ae9c016b71f402e43208e13ddf00912f47f623ee5994e12fc7d4b3f \
+    --hash=sha256:c02a30b904282777d872266b87b20ed8cc0d1501855e27f831320f471d54d983 \
+    --hash=sha256:c1277cd707c465cd09572a774559a3cc7c7a28802eb3a2a9472588f062097205 \
+    --hash=sha256:c30a0eafc89d28e7f959281b58198a9fa5e99405f716c0289b7892ca345fe45f \
+    --hash=sha256:c5ce2ce7c997e1971b7184ee37deb6ea9922ef5163c6ee5aa3c274b05f9e12fa \
+    --hash=sha256:c823bc3971c44ab93e611ab1a46b1eafeae474c0c844aff4b7474287b75fe49c \
+    --hash=sha256:ce0cdc074d540265bfeb31336e678b4e37316849d13b308607efa527e981f5c2 \
+    --hash=sha256:d1720b4f14c78a3089562b8875b53e36b51c97c51adc53325a69b79b4b48ebcb \
+    --hash=sha256:d183cf9c797a5291e8301790ed6d053480ed94070637bfaad914dd38b0981f67 \
+    --hash=sha256:d9010c31cd6fa59438da4e58a7f19e4753f7f264300cd152e7f90d4602449762 \
+    --hash=sha256:d9e5e4a85bdb56d224f412d9c98ae4cbd032cc4f3161818f692cd81766eee65a \
+    --hash=sha256:da1dee8948d2137bb51fbb8a53cce6b1bcc86003c6b42565f008438b806cccd8 \
+    --hash=sha256:df9270660711670e68803107d55c2b5949c2e0f2e4896da176e1ecfc068b974a \
+    --hash=sha256:e00e3505cd80440f6c98c6d69269dcc2a119f86ad0a9fd70bccc59504bebd68a \
+    --hash=sha256:e48d5021a84d341bcaf95c8460b152cfbad770d28e5fe14a768988c461b821bc \
+    --hash=sha256:e7f8b04d83483577fd9200461b057c9f14ced334dcb053090cea1da9c8321a91 \
+    --hash=sha256:edfe3341033a6b53a5c522c802deb2079eee5cbfbb0af032a55064bd65c73a23 \
+    --hash=sha256:ef9c33cc5cbca35808f6c74be11eb7f5f6b14d2311be84a15b594bd3e58b5527 \
+    --hash=sha256:f2d4324a98062be0525d16f768a03e0bbb3b9fe301ceee99611dc9a7953124e6 \
+    --hash=sha256:f3935f82f6f4a3820270842e90456ebad3af15810cf65932bd24da4463bc0a4c \
+    --hash=sha256:f614ab0c76397661b90b6851a030004dac502e48260ea10f2441abd2207fbcc7 \
+    --hash=sha256:f7db54c7914cc99d901d93a34704833568d86c20925b2762f9fa779f9cd2e70f \
+    --hash=sha256:fbc6264158392bad9df19537e872d476f7c57adf718944cc1e4495cbabf38e2a \
+    --hash=sha256:fe2fb38c2ed905a2582948e2de560675e9dfbee94c6d5ccdb1301c6d0a5bf092 \
+    --hash=sha256:ffe595f10566f8276b76dc3a11ae4bb7eba1aac8ddd75811736a15b0d5311414
+    # via gcsfs
+aiosignal==1.3.1 \
+    --hash=sha256:54cd96e15e1649b75d6c87526a6ff0b6c1b0dd3459f43d9ca11d48c339b68cfc \
+    --hash=sha256:f8376fb07dd1e86a584e4fcdec80b36b7f81aac666ebc724e2c090300dd83b17
+    # via aiohttp
+asciitree==0.3.3 \
+    --hash=sha256:4aa4b9b649f85e3fcb343363d97564aa1fb62e249677f2e18a96765145cc0f6e
+    # via zarr
 attrs==24.2.0 \
     --hash=sha256:5cfb1b9148b5b086569baec03f20d7b6bf3bcacc9a42bebf87ffaaca362f6346 \
     --hash=sha256:81921eb96de3191c8258c199618104dd27ac608d9366f5e35d011eae1867ede2
-    # via e3nn-jax
+    # via
+    #   aiohttp
+    #   e3nn-jax
+cachetools==5.5.0 \
+    --hash=sha256:02134e8439cdc2ffb62023ce1debca2944c3f289d66bb17ead3ab3dede74b292 \
+    --hash=sha256:2cc24fb4cbe39633fb7badd9db9ca6295d766d9c2995f245725a46715d050f2a
+    # via google-auth
 certifi==2024.7.4 \
     --hash=sha256:5a1e7645bc0ec61a09e26c36f6106dd4cf40c6db3a1fb6352b0244e7fb057c7b \
     --hash=sha256:c198e21b1289c2ab85ee4e67bb4b4ef3ead0892059901a8d5b622f24a1101e90
@@ -120,14 +231,78 @@ chex==0.1.87 \
     --hash=sha256:0096d89cc8d898bb521ef4bfbf5c24549022b0e5b301f529ab57238896fe6c5d \
     --hash=sha256:ce536475661fd96d21be0c1728ecdbedd03f8ff950c662dfc338c92ea782cb16
     # via optax
+cloudpickle==3.1.0 \
+    --hash=sha256:81a929b6e3c7335c863c771d673d105f02efdb89dfaba0c90495d1c64796601b \
+    --hash=sha256:fe11acda67f61aaaec473e3afe030feb131d78a43461b718185363384f1ba12e
+    # via tensorflow-probability
 contextlib2==21.6.0 \
     --hash=sha256:3fbdb64466afd23abaf6c977627b75b6139a5a3e8ce38405c5b413aed7a0471f \
     --hash=sha256:ab1e2bfe1d01d968e1b7e8d9023bc51ef3509bba217bb730cee3827e1ee82869
     # via ml-collections
+decorator==5.1.1 \
+    --hash=sha256:637996211036b6385ef91435e4fae22989472f9d571faba8927ba8253acbc330 \
+    --hash=sha256:b8c3f85900b9dc423225913c5aace94729fe1fa9763b38939a95226f02d37186
+    # via
+    #   gcsfs
+    #   tensorflow-probability
+dinosaur-dycore==1.0.0 \
+    --hash=sha256:0ea3807a83d0c5220ff7a1cd548c11d82091d455d17d8f7697b985b6283d2ef9 \
+    --hash=sha256:3505d59eb9d90ad58572f93fa87db5536b9df0663f95701d50c8ccb027e767e1
+    # via neuralgcm
 dm-haiku==0.0.12 \
     --hash=sha256:7448a43a6486bff95253f84e18eacc607d9c1256592573117a9d1d23e2780706 \
     --hash=sha256:ba0b3acf71433156737fe342c486da11727e5e6c9e054245f4f9b8f0b53eb608
-    # via jax-md
+    # via
+    #   jax-md
+    #   neuralgcm
+dm-tree==0.1.8 \
+    --hash=sha256:054b461f8176f4bce7a21f7b1870f873a1ced3bdbe1282c816c550bb43c71fa6 \
+    --hash=sha256:09964470f76a5201aff2e8f9b26842976de7889300676f927930f6285e256760 \
+    --hash=sha256:0d3172394079a86c3a759179c65f64c48d1a42b89495fcf38976d11cc3bb952c \
+    --hash=sha256:0e9620ccf06393eb6b613b5e366469304622d4ea96ae6540b28a33840e6c89cf \
+    --hash=sha256:0fcaabbb14e7980377439e7140bd05552739ca5e515ecb3119f234acee4b9430 \
+    --hash=sha256:1607ce49aa42f010d1e5e616d92ce899d66835d4d8bea49679582435285515de \
+    --hash=sha256:181c35521d480d0365f39300542cb6cd7fd2b77351bb43d7acfda15aef63b317 \
+    --hash=sha256:1d7c26e431fc93cc7e0cba867eb000db6a05f6f2b25af11ac4e9dada88fc5bca \
+    --hash=sha256:1fe962015b2fe1282892b28ebe962faed53c7f98d942da9a4625cbf27baef913 \
+    --hash=sha256:250b692fb75f45f02e2f58fbef9ab338904ef334b90557565621fa251df267cf \
+    --hash=sha256:2869228d9c619074de501a3c10dc7f07c75422f8fab36ecdcb859b6f1b1ec3ef \
+    --hash=sha256:28c52cbf4f8b3dbd0beaedf44f69fa85eec5e9dede612e08035e06ada6ec9426 \
+    --hash=sha256:2f7915660f59c09068e428613c480150180df1060561fd0d1470684ae7007bd1 \
+    --hash=sha256:343a4a4ebaa127451ff971254a4be4084eb4bdc0b2513c32b46f6f728fd03f9e \
+    --hash=sha256:35cc164a79336bfcfafb47e5f297898359123bbd3330c1967f0c4994f9cf9f60 \
+    --hash=sha256:378cc8ad93c5fe3590f405a309980721f021c790ca1bdf9b15bb1d59daec57f5 \
+    --hash=sha256:39070ba268c0491af9fe7a58644d99e8b4f2cde6e5884ba3380bddc84ed43d5f \
+    --hash=sha256:435227cf3c5dc63f4de054cf3d00183790bd9ead4c3623138c74dde7f67f521b \
+    --hash=sha256:5483dca4d7eb1a0d65fe86d3b6a53ae717face83c1f17e0887b1a4a64ae5c410 \
+    --hash=sha256:694c3654cfd2a81552c08ec66bb5c4a3d48fa292b9a181880fb081c36c5b9134 \
+    --hash=sha256:75c5d528bb992981c20793b6b453e91560784215dffb8a5440ba999753c14ceb \
+    --hash=sha256:803bfc53b4659f447ac694dbd04235f94a73ef7c1fd1e0df7c84ac41e0bc963b \
+    --hash=sha256:81fce77f22a302d7a5968aebdf4efafef4def7ce96528719a354e6990dcd49c7 \
+    --hash=sha256:83b7764de0d855338abefc6e3ee9fe40d301668310aa3baea3f778ff051f4393 \
+    --hash=sha256:8c60a7eadab64c2278861f56bca320b2720f163dca9d7558103c3b77f2416571 \
+    --hash=sha256:8ed3564abed97c806db122c2d3e1a2b64c74a63debe9903aad795167cc301368 \
+    --hash=sha256:94d3f0826311f45ee19b75f5b48c99466e4218a0489e81c0f0167bda50cacf22 \
+    --hash=sha256:96a548a406a6fb15fe58f6a30a57ff2f2aafbf25f05afab00c8f5e5977b6c715 \
+    --hash=sha256:a5d819c38c03f0bb5b3b3703c60e4b170355a0fc6b5819325bf3d4ceb3ae7e80 \
+    --hash=sha256:ad16ceba90a56ec47cf45b21856d14962ac314787975ef786efb5e6e9ca75ec7 \
+    --hash=sha256:af4b3d372f2477dcd89a6e717e4a575ca35ccc20cc4454a8a4b6f8838a00672d \
+    --hash=sha256:b095ba4f8ca1ba19350fd53cf1f8f3eb0bd406aa28af64a6dfc86707b32a810a \
+    --hash=sha256:b9bd9b9ccb59409d33d51d84b7668010c04c2af7d4a371632874c1ca356cff3d \
+    --hash=sha256:b9f89a454e98806b44fe9d40ec9eee61f848388f7e79ac2371a55679bd5a3ac6 \
+    --hash=sha256:bb2d109f42190225112da899b9f3d46d0d5f26aef501c61e43529fe9322530b5 \
+    --hash=sha256:c0a94aba18a35457a1b5cd716fd7b46c5dafdc4cf7869b4bae665b91c4682a8e \
+    --hash=sha256:c5c8c12e3fda754ef6af94161bacdaeda816d941995fac415d6855c6c386af68 \
+    --hash=sha256:d1612fcaecd79023dbc6a6ae48d51a80beb5c385d6f3f6d71688e57bc8d07de8 \
+    --hash=sha256:d16e1f2a073604cfcc09f7131ae8d534674f43c3aef4c25742eae295bc60d04f \
+    --hash=sha256:d20f2faa3672b52e5013f4077117bfb99c4cfc0b445d3bde1584c34032b57436 \
+    --hash=sha256:d40fa4106ca6edc66760246a08f500ec0c85ef55c762fb4a363f6ee739ba02ee \
+    --hash=sha256:de287fabc464b8734be251e46e06aa9aa1001f34198da2b6ce07bd197172b9cb \
+    --hash=sha256:e4d714371bb08839e4e5e29024fc95832d9affe129825ef38836b143028bd144 \
+    --hash=sha256:ea9e59e0451e7d29aece402d9f908f2e2a80922bcde2ebfd5dcb07750fcbfee8 \
+    --hash=sha256:f7ac31b9aecccb2c6e1ab29706f6ded3eba0c2c69c770322c9c685929c3d6afb \
+    --hash=sha256:fa42a605d099ee7d41ba2b5fb75e21423951fd26e5d50583a00471238fb3021d
+    # via tensorflow-probability
 e3nn-jax==0.20.7 \
     --hash=sha256:7764edb99d19465e574a23857e5f97420b4d720f06ba8bc18a01f9f2b8301d1c \
     --hash=sha256:8a1b6f0c00feed565358fceec4c916a623412b6328945fce272f0c33e52b4bfc
@@ -142,12 +317,24 @@ etils[epath,epy]==1.9.4 \
     # via
     #   optax
     #   orbax-checkpoint
+fasteners==0.19 \
+    --hash=sha256:758819cb5d94cdedf4e836988b74de396ceacb8e2794d21f82d131fd9ee77237 \
+    --hash=sha256:b4f37c3ac52d8a445af3a66bce57b33b5e90b97c696b7b984f530cf8f0ded09c
+    # via zarr
 flax==0.9.0 \
     --hash=sha256:12cd8f7162165ddd56877fb1cd9a4fcb47a31569e4c5343eeb59a36369fa2cfe \
     --hash=sha256:8b7f361eed0f5324e81f9dc8d02ea53da5f993d7c2e37e7aa5b37d3f6331dd53
     # via
     #   dm-haiku
     #   jax-md
+flexcache==0.3 \
+    --hash=sha256:18743bd5a0621bfe2cf8d519e4c3bfdf57a269c15d1ced3fb4b64e0ff4600656 \
+    --hash=sha256:d43c9fea82336af6e0115e308d9d33a185390b8346a017564611f1466dcd2e32
+    # via pint
+flexparser==0.4 \
+    --hash=sha256:266d98905595be2ccc5da964fe0a2c3526fbbffdc45b65b3146d75db992ef6b2 \
+    --hash=sha256:3738b456192dcb3e15620f324c447721023c0293f6af9955b481e91d00179846
+    # via pint
 frozendict==2.4.4 \
     --hash=sha256:07c3a5dee8bbb84cba770e273cdbf2c87c8e035903af8f781292d72583416801 \
     --hash=sha256:12a342e439aef28ccec533f0253ea53d75fe9102bd6ea928ff530e76eac38906 \
@@ -183,10 +370,187 @@ frozendict==2.4.4 \
     --hash=sha256:e1b941132d79ce72d562a13341d38fc217bc1ee24d8c35a20d754e79ff99e038 \
     --hash=sha256:fefeb700bc7eb8b4c2dc48704e4221860d254c8989fb53488540bc44e44a1ac2
     # via jax-md
+frozenlist==1.5.0 \
+    --hash=sha256:000a77d6034fbad9b6bb880f7ec073027908f1b40254b5d6f26210d2dab1240e \
+    --hash=sha256:03d33c2ddbc1816237a67f66336616416e2bbb6beb306e5f890f2eb22b959cdf \
+    --hash=sha256:04a5c6babd5e8fb7d3c871dc8b321166b80e41b637c31a995ed844a6139942b6 \
+    --hash=sha256:0996c66760924da6e88922756d99b47512a71cfd45215f3570bf1e0b694c206a \
+    --hash=sha256:0cc974cc93d32c42e7b0f6cf242a6bd941c57c61b618e78b6c0a96cb72788c1d \
+    --hash=sha256:0f253985bb515ecd89629db13cb58d702035ecd8cfbca7d7a7e29a0e6d39af5f \
+    --hash=sha256:11aabdd62b8b9c4b84081a3c246506d1cddd2dd93ff0ad53ede5defec7886b28 \
+    --hash=sha256:12f78f98c2f1c2429d42e6a485f433722b0061d5c0b0139efa64f396efb5886b \
+    --hash=sha256:140228863501b44b809fb39ec56b5d4071f4d0aa6d216c19cbb08b8c5a7eadb9 \
+    --hash=sha256:1431d60b36d15cda188ea222033eec8e0eab488f39a272461f2e6d9e1a8e63c2 \
+    --hash=sha256:15538c0cbf0e4fa11d1e3a71f823524b0c46299aed6e10ebb4c2089abd8c3bec \
+    --hash=sha256:15b731db116ab3aedec558573c1a5eec78822b32292fe4f2f0345b7f697745c2 \
+    --hash=sha256:17dcc32fc7bda7ce5875435003220a457bcfa34ab7924a49a1c19f55b6ee185c \
+    --hash=sha256:1893f948bf6681733aaccf36c5232c231e3b5166d607c5fa77773611df6dc336 \
+    --hash=sha256:189f03b53e64144f90990d29a27ec4f7997d91ed3d01b51fa39d2dbe77540fd4 \
+    --hash=sha256:1a8ea951bbb6cacd492e3948b8da8c502a3f814f5d20935aae74b5df2b19cf3d \
+    --hash=sha256:1b96af8c582b94d381a1c1f51ffaedeb77c821c690ea5f01da3d70a487dd0a9b \
+    --hash=sha256:1e76bfbc72353269c44e0bc2cfe171900fbf7f722ad74c9a7b638052afe6a00c \
+    --hash=sha256:2150cc6305a2c2ab33299453e2968611dacb970d2283a14955923062c8d00b10 \
+    --hash=sha256:226d72559fa19babe2ccd920273e767c96a49b9d3d38badd7c91a0fdeda8ea08 \
+    --hash=sha256:237f6b23ee0f44066219dae14c70ae38a63f0440ce6750f868ee08775073f942 \
+    --hash=sha256:29d94c256679247b33a3dc96cce0f93cbc69c23bf75ff715919332fdbb6a32b8 \
+    --hash=sha256:2b5e23253bb709ef57a8e95e6ae48daa9ac5f265637529e4ce6b003a37b2621f \
+    --hash=sha256:2d0da8bbec082bf6bf18345b180958775363588678f64998c2b7609e34719b10 \
+    --hash=sha256:2f3f7a0fbc219fb4455264cae4d9f01ad41ae6ee8524500f381de64ffaa077d5 \
+    --hash=sha256:30c72000fbcc35b129cb09956836c7d7abf78ab5416595e4857d1cae8d6251a6 \
+    --hash=sha256:31115ba75889723431aa9a4e77d5f398f5cf976eea3bdf61749731f62d4a4a21 \
+    --hash=sha256:31a9ac2b38ab9b5a8933b693db4939764ad3f299fcaa931a3e605bc3460e693c \
+    --hash=sha256:366d8f93e3edfe5a918c874702f78faac300209a4d5bf38352b2c1bdc07a766d \
+    --hash=sha256:374ca2dabdccad8e2a76d40b1d037f5bd16824933bf7bcea3e59c891fd4a0923 \
+    --hash=sha256:44c49271a937625619e862baacbd037a7ef86dd1ee215afc298a417ff3270608 \
+    --hash=sha256:45e0896250900b5aa25180f9aec243e84e92ac84bd4a74d9ad4138ef3f5c97de \
+    --hash=sha256:498524025a5b8ba81695761d78c8dd7382ac0b052f34e66939c42df860b8ff17 \
+    --hash=sha256:50cf5e7ee9b98f22bdecbabf3800ae78ddcc26e4a435515fc72d97903e8488e0 \
+    --hash=sha256:52ef692a4bc60a6dd57f507429636c2af8b6046db8b31b18dac02cbc8f507f7f \
+    --hash=sha256:561eb1c9579d495fddb6da8959fd2a1fca2c6d060d4113f5844b433fc02f2641 \
+    --hash=sha256:5a3ba5f9a0dfed20337d3e966dc359784c9f96503674c2faf015f7fe8e96798c \
+    --hash=sha256:5b6a66c18b5b9dd261ca98dffcb826a525334b2f29e7caa54e182255c5f6a65a \
+    --hash=sha256:5c28f4b5dbef8a0d8aad0d4de24d1e9e981728628afaf4ea0792f5d0939372f0 \
+    --hash=sha256:5d7f5a50342475962eb18b740f3beecc685a15b52c91f7d975257e13e029eca9 \
+    --hash=sha256:6321899477db90bdeb9299ac3627a6a53c7399c8cd58d25da094007402b039ab \
+    --hash=sha256:6482a5851f5d72767fbd0e507e80737f9c8646ae7fd303def99bfe813f76cf7f \
+    --hash=sha256:666534d15ba8f0fda3f53969117383d5dc021266b3c1a42c9ec4855e4b58b9d3 \
+    --hash=sha256:683173d371daad49cffb8309779e886e59c2f369430ad28fe715f66d08d4ab1a \
+    --hash=sha256:6e9080bb2fb195a046e5177f10d9d82b8a204c0736a97a153c2466127de87784 \
+    --hash=sha256:73f2e31ea8dd7df61a359b731716018c2be196e5bb3b74ddba107f694fbd7604 \
+    --hash=sha256:7437601c4d89d070eac8323f121fcf25f88674627505334654fd027b091db09d \
+    --hash=sha256:76e4753701248476e6286f2ef492af900ea67d9706a0155335a40ea21bf3b2f5 \
+    --hash=sha256:7707a25d6a77f5d27ea7dc7d1fc608aa0a478193823f88511ef5e6b8a48f9d03 \
+    --hash=sha256:7948140d9f8ece1745be806f2bfdf390127cf1a763b925c4a805c603df5e697e \
+    --hash=sha256:7a1a048f9215c90973402e26c01d1cff8a209e1f1b53f72b95c13db61b00f953 \
+    --hash=sha256:7d57d8f702221405a9d9b40f9da8ac2e4a1a8b5285aac6100f3393675f0a85ee \
+    --hash=sha256:7f3c8c1dacd037df16e85227bac13cca58c30da836c6f936ba1df0c05d046d8d \
+    --hash=sha256:81d5af29e61b9c8348e876d442253723928dce6433e0e76cd925cd83f1b4b817 \
+    --hash=sha256:828afae9f17e6de596825cf4228ff28fbdf6065974e5ac1410cecc22f699d2b3 \
+    --hash=sha256:87f724d055eb4785d9be84e9ebf0f24e392ddfad00b3fe036e43f489fafc9039 \
+    --hash=sha256:8969190d709e7c48ea386db202d708eb94bdb29207a1f269bab1196ce0dcca1f \
+    --hash=sha256:90646abbc7a5d5c7c19461d2e3eeb76eb0b204919e6ece342feb6032c9325ae9 \
+    --hash=sha256:91d6c171862df0a6c61479d9724f22efb6109111017c87567cfeb7b5d1449fdf \
+    --hash=sha256:9272fa73ca71266702c4c3e2d4a28553ea03418e591e377a03b8e3659d94fa76 \
+    --hash=sha256:92b5278ed9d50fe610185ecd23c55d8b307d75ca18e94c0e7de328089ac5dcba \
+    --hash=sha256:97160e245ea33d8609cd2b8fd997c850b56db147a304a262abc2b3be021a9171 \
+    --hash=sha256:977701c081c0241d0955c9586ffdd9ce44f7a7795df39b9151cd9a6fd0ce4cfb \
+    --hash=sha256:9b7dc0c4338e6b8b091e8faf0db3168a37101943e687f373dce00959583f7439 \
+    --hash=sha256:9b93d7aaa36c966fa42efcaf716e6b3900438632a626fb09c049f6a2f09fc631 \
+    --hash=sha256:9bbcdfaf4af7ce002694a4e10a0159d5a8d20056a12b05b45cea944a4953f972 \
+    --hash=sha256:9c2623347b933fcb9095841f1cc5d4ff0b278addd743e0e966cb3d460278840d \
+    --hash=sha256:a2fe128eb4edeabe11896cb6af88fca5346059f6c8d807e3b910069f39157869 \
+    --hash=sha256:a72b7a6e3cd2725eff67cd64c8f13335ee18fc3c7befc05aed043d24c7b9ccb9 \
+    --hash=sha256:a9fe0f1c29ba24ba6ff6abf688cb0b7cf1efab6b6aa6adc55441773c252f7411 \
+    --hash=sha256:b97f7b575ab4a8af9b7bc1d2ef7f29d3afee2226bd03ca3875c16451ad5a7723 \
+    --hash=sha256:bdac3c7d9b705d253b2ce370fde941836a5f8b3c5c2b8fd70940a3ea3af7f4f2 \
+    --hash=sha256:c03eff4a41bd4e38415cbed054bbaff4a075b093e2394b6915dca34a40d1e38b \
+    --hash=sha256:c16d2fa63e0800723139137d667e1056bee1a1cf7965153d2d104b62855e9b99 \
+    --hash=sha256:c1fac3e2ace2eb1052e9f7c7db480818371134410e1f5c55d65e8f3ac6d1407e \
+    --hash=sha256:ce3aa154c452d2467487765e3adc730a8c153af77ad84096bc19ce19a2400840 \
+    --hash=sha256:cee6798eaf8b1416ef6909b06f7dc04b60755206bddc599f52232606e18179d3 \
+    --hash=sha256:d1b3eb7b05ea246510b43a7e53ed1653e55c2121019a97e60cad7efb881a97bb \
+    --hash=sha256:d994863bba198a4a518b467bb971c56e1db3f180a25c6cf7bb1949c267f748c3 \
+    --hash=sha256:dd47a5181ce5fcb463b5d9e17ecfdb02b678cca31280639255ce9d0e5aa67af0 \
+    --hash=sha256:dd94994fc91a6177bfaafd7d9fd951bc8689b0a98168aa26b5f543868548d3ca \
+    --hash=sha256:de537c11e4aa01d37db0d403b57bd6f0546e71a82347a97c6a9f0dcc532b3a45 \
+    --hash=sha256:df6e2f325bfee1f49f81aaac97d2aa757c7646534a06f8f577ce184afe2f0a9e \
+    --hash=sha256:e66cc454f97053b79c2ab09c17fbe3c825ea6b4de20baf1be28919460dd7877f \
+    --hash=sha256:e79225373c317ff1e35f210dd5f1344ff31066ba8067c307ab60254cd3a78ad5 \
+    --hash=sha256:f1577515d35ed5649d52ab4319db757bb881ce3b2b796d7283e6634d99ace307 \
+    --hash=sha256:f1e6540b7fa044eee0bb5111ada694cf3dc15f2b0347ca125ee9ca984d5e9e6e \
+    --hash=sha256:f2ac49a9bedb996086057b75bf93538240538c6d9b38e57c82d51f75a73409d2 \
+    --hash=sha256:f47c9c9028f55a04ac254346e92977bf0f166c483c74b4232bee19a6697e4778 \
+    --hash=sha256:f5f9da7f5dbc00a604fe74aa02ae7c98bcede8a3b8b9666f9f86fc13993bc71a \
+    --hash=sha256:fd74520371c3c4175142d02a976aee0b4cb4a7cc912a60586ffd8d5929979b30 \
+    --hash=sha256:feeb64bc9bcc6b45c6311c9e9b99406660a9c05ca8a5b30d14a78555088b0b3a
+    # via
+    #   aiohttp
+    #   aiosignal
 fsspec==2024.9.0 \
     --hash=sha256:4b0afb90c2f21832df142f292649035d80b421f60a9e1c027802e5a0da2b04e8 \
     --hash=sha256:a0947d552d8a6efa72cc2c730b12c41d043509156966cca4fb157b0f2a0c574b
-    # via etils
+    # via
+    #   dinosaur-dycore
+    #   etils
+    #   gcsfs
+gast==0.6.0 \
+    --hash=sha256:52b182313f7330389f72b069ba00f174cfe2a06411099547288839c6cbafbd54 \
+    --hash=sha256:88fc5300d32c7ac6ca7b515310862f71e6fdf2c029bbec7c66c0f5dd47b6b1fb
+    # via tensorflow-probability
+gcsfs==2024.9.0.post1 \
+    --hash=sha256:7ca70ee9d7c7dbce1a3e36b4883e14102c2d7b4284f49e242843a437bc684684 \
+    --hash=sha256:f3ab9d3bedc45da8cf40baed7c3a1e1694e8f599160d9138d78f0ef25e4a3ca1
+    # via -r builddeps/test-requirements.txt
+gin-config==0.5.0 \
+    --hash=sha256:0c6ea5026ded927c8c93c990b01c695257c1df446e45e549a158cfbc79e19ed6 \
+    --hash=sha256:bddb7ca221ea2b46cdb59321e79fecf02d6e3b728906047fcd4076c297609fd6
+    # via neuralgcm
+google-api-core==2.22.0 \
+    --hash=sha256:26f8d76b96477db42b55fd02a33aae4a42ec8b86b98b94969b7333a2c828bf35 \
+    --hash=sha256:a6652b6bd51303902494998626653671703c420f6f4c88cfd3f50ed723e9d021
+    # via
+    #   google-cloud-core
+    #   google-cloud-storage
+google-auth==2.36.0 \
+    --hash=sha256:51a15d47028b66fd36e5c64a82d2d57480075bccc7da37cde257fc94177a61fb \
+    --hash=sha256:545e9618f2df0bcbb7dcbc45a546485b1212624716975a1ea5ae8149ce769ab1
+    # via
+    #   gcsfs
+    #   google-api-core
+    #   google-auth-oauthlib
+    #   google-cloud-core
+    #   google-cloud-storage
+google-auth-oauthlib==1.2.1 \
+    --hash=sha256:2d58a27262d55aa1b87678c3ba7142a080098cbc2024f903c62355deb235d91f \
+    --hash=sha256:afd0cad092a2eaa53cd8e8298557d6de1034c6cb4a740500b5357b648af97263
+    # via gcsfs
+google-cloud-core==2.4.1 \
+    --hash=sha256:9b7749272a812bde58fff28868d0c5e2f585b82f37e09a1f6ed2d4d10f134073 \
+    --hash=sha256:a9e6a4422b9ac5c29f79a0ede9485473338e2ce78d91f2370c01e730eab22e61
+    # via google-cloud-storage
+google-cloud-storage==2.18.2 \
+    --hash=sha256:97a4d45c368b7d401ed48c4fdfe86e1e1cb96401c9e199e419d289e2c0370166 \
+    --hash=sha256:aaf7acd70cdad9f274d29332673fcab98708d0e1f4dceb5a5356aaef06af4d99
+    # via gcsfs
+google-crc32c==1.6.0 \
+    --hash=sha256:05e2d8c9a2f853ff116db9706b4a27350587f341eda835f46db3c0a8c8ce2f24 \
+    --hash=sha256:18e311c64008f1f1379158158bb3f0c8d72635b9eb4f9545f8cf990c5668e59d \
+    --hash=sha256:236c87a46cdf06384f614e9092b82c05f81bd34b80248021f729396a78e55d7e \
+    --hash=sha256:35834855408429cecf495cac67ccbab802de269e948e27478b1e47dfb6465e57 \
+    --hash=sha256:386122eeaaa76951a8196310432c5b0ef3b53590ef4c317ec7588ec554fec5d2 \
+    --hash=sha256:40b05ab32a5067525670880eb5d169529089a26fe35dce8891127aeddc1950e8 \
+    --hash=sha256:48abd62ca76a2cbe034542ed1b6aee851b6f28aaca4e6551b5599b6f3ef175cc \
+    --hash=sha256:50cf2a96da226dcbff8671233ecf37bf6e95de98b2a2ebadbfdf455e6d05df42 \
+    --hash=sha256:51c4f54dd8c6dfeb58d1df5e4f7f97df8abf17a36626a217f169893d1d7f3e9f \
+    --hash=sha256:5bcc90b34df28a4b38653c36bb5ada35671ad105c99cfe915fb5bed7ad6924aa \
+    --hash=sha256:62f6d4a29fea082ac4a3c9be5e415218255cf11684ac6ef5488eea0c9132689b \
+    --hash=sha256:6eceb6ad197656a1ff49ebfbbfa870678c75be4344feb35ac1edf694309413dc \
+    --hash=sha256:7aec8e88a3583515f9e0957fe4f5f6d8d4997e36d0f61624e70469771584c760 \
+    --hash=sha256:91ca8145b060679ec9176e6de4f89b07363d6805bd4760631ef254905503598d \
+    --hash=sha256:a184243544811e4a50d345838a883733461e67578959ac59964e43cca2c791e7 \
+    --hash=sha256:a9e4b426c3702f3cd23b933436487eb34e01e00327fac20c9aebb68ccf34117d \
+    --hash=sha256:bb0966e1c50d0ef5bc743312cc730b533491d60585a9a08f897274e57c3f70e0 \
+    --hash=sha256:bb8b3c75bd157010459b15222c3fd30577042a7060e29d42dabce449c087f2b3 \
+    --hash=sha256:bd5e7d2445d1a958c266bfa5d04c39932dc54093fa391736dbfdb0f1929c1fb3 \
+    --hash=sha256:c87d98c7c4a69066fd31701c4e10d178a648c2cac3452e62c6b24dc51f9fcc00 \
+    --hash=sha256:d2952396dc604544ea7476b33fe87faedc24d666fb0c2d5ac971a2b9576ab871 \
+    --hash=sha256:d8797406499f28b5ef791f339594b0b5fdedf54e203b5066675c406ba69d705c \
+    --hash=sha256:d9e9913f7bd69e093b81da4535ce27af842e7bf371cde42d1ae9e9bd382dc0e9 \
+    --hash=sha256:e2806553238cd076f0a55bddab37a532b53580e699ed8e5606d0de1f856b5205 \
+    --hash=sha256:ebab974b1687509e5c973b5c4b8b146683e101e102e17a86bd196ecaa4d099fc \
+    --hash=sha256:ed767bf4ba90104c1216b68111613f0d5926fb3780660ea1198fc469af410e9d \
+    --hash=sha256:f7a1fc29803712f80879b0806cb83ab24ce62fc8daf0569f2204a0cfd7f68ed4
+    # via
+    #   google-cloud-storage
+    #   google-resumable-media
+google-resumable-media==2.7.2 \
+    --hash=sha256:3ce7551e9fe6d99e9a126101d2536612bb73486721951e9562fee0f90c6ababa \
+    --hash=sha256:5280aed4629f2b60b847b0d42f9857fd4935c11af266744df33d8074cae92fe0
+    # via google-cloud-storage
+googleapis-common-protos==1.65.0 \
+    --hash=sha256:2972e6c496f435b92590fd54045060867f3fe9be2c82ab148fc8885035479a63 \
+    --hash=sha256:334a29d07cddc3aa01dee4988f9afd9b2916ee2ff49d6b757155dc0d197852c0
+    # via google-api-core
 humanize==4.11.0 \
     --hash=sha256:b53caaec8532bcb2fff70c8826f904c35943f8cecaca29d272d9df38092736c0 \
     --hash=sha256:e66f36020a2d5a974c504bd2555cf770621dbdbb6d82f94a6857c0b1ea2608be
@@ -194,7 +558,9 @@ humanize==4.11.0 \
 idna==3.8 \
     --hash=sha256:050b4e5baadcd44d760cedbd2b8e639f2ff89bbc7a5730fcc662954303377aac \
     --hash=sha256:d838c2c0ed6fced7693d5e8ab8e734d5f8fda53a039c0164afb0b82e771e3603
-    # via requests
+    # via
+    #   requests
+    #   yarl
 importlib-resources==6.4.5 \
     --hash=sha256:980862a1d16c9e147a59603677fa2aa5fd82b87f223b6cb870695bcfce830065 \
     --hash=sha256:ac29d5f956f01d5e4bb63102a5a19957f1b9175e45649977264a1416783bb717
@@ -206,12 +572,15 @@ jax==0.4.31 \
     #   -r builddeps/requirements.in
     #   -r builddeps/test-requirements.txt
     #   chex
+    #   dinosaur-dycore
     #   e3nn-jax
     #   flax
     #   jax-md
     #   jraph
+    #   neuralgcm
     #   optax
     #   orbax-checkpoint
+    #   tree-math
 jax-cuda12-pjrt==0.4.31 \
     --hash=sha256:3e77d1cfebeca06517254eb568f082037e1a2aa3ed8f63c543492ad8ab5a1585 \
     --hash=sha256:8961abb381d893a3c2392ad76ab2067a81f8f2514f3b47d2da3ac24283293fe0
@@ -247,16 +616,22 @@ jaxlib==0.4.31 \
     #   -r builddeps/requirements.in
     #   -r builddeps/test-requirements.txt
     #   chex
+    #   dinosaur-dycore
     #   e3nn-jax
     #   jax
     #   jax-md
     #   jraph
+    #   neuralgcm
     #   optax
     #   orbax-checkpoint
 jmp==0.0.4 \
     --hash=sha256:5dfeb0fd7c7a9f72a70fff0aab9d0cbfae32a809c02f4037ff3485ceb33e1730 \
     --hash=sha256:6aa7adbddf2bd574b28c7faf6e81a735eb11f53386447896909c6968dc36807d
     # via dm-haiku
+joblib==1.4.2 \
+    --hash=sha256:06d478d5674cbc267e7496a410ee875abd68e4340feff4490bcb7afb88060ae6 \
+    --hash=sha256:2382c5816b2636fbd20a09e0f4e9dad4736765fdfb7dca582943b9c1366b3f0e
+    # via scikit-learn
 jraph @ https://github.com/wsmoses/jraph/archive/b00d9a03db76c69a258a86df81638b9a2f28829f.tar.gz \
     --hash=sha256:d996f47fff8c6ac8f5572e4220e3193f2542e9d3d6bc1defc16ac22dadd74dd4
     # via -r builddeps/test-requirements.txt
@@ -365,10 +740,129 @@ msgpack==1.1.0 \
     # via
     #   flax
     #   orbax-checkpoint
+multidict==6.1.0 \
+    --hash=sha256:052e10d2d37810b99cc170b785945421141bf7bb7d2f8799d431e7db229c385f \
+    --hash=sha256:06809f4f0f7ab7ea2cabf9caca7d79c22c0758b58a71f9d32943ae13c7ace056 \
+    --hash=sha256:071120490b47aa997cca00666923a83f02c7fbb44f71cf7f136df753f7fa8761 \
+    --hash=sha256:0c3f390dc53279cbc8ba976e5f8035eab997829066756d811616b652b00a23a3 \
+    --hash=sha256:0e2b90b43e696f25c62656389d32236e049568b39320e2735d51f08fd362761b \
+    --hash=sha256:0e5f362e895bc5b9e67fe6e4ded2492d8124bdf817827f33c5b46c2fe3ffaca6 \
+    --hash=sha256:10524ebd769727ac77ef2278390fb0068d83f3acb7773792a5080f2b0abf7748 \
+    --hash=sha256:10a9b09aba0c5b48c53761b7c720aaaf7cf236d5fe394cd399c7ba662d5f9966 \
+    --hash=sha256:16e5f4bf4e603eb1fdd5d8180f1a25f30056f22e55ce51fb3d6ad4ab29f7d96f \
+    --hash=sha256:188215fc0aafb8e03341995e7c4797860181562380f81ed0a87ff455b70bf1f1 \
+    --hash=sha256:189f652a87e876098bbc67b4da1049afb5f5dfbaa310dd67c594b01c10388db6 \
+    --hash=sha256:1ca0083e80e791cffc6efce7660ad24af66c8d4079d2a750b29001b53ff59ada \
+    --hash=sha256:1e16bf3e5fc9f44632affb159d30a437bfe286ce9e02754759be5536b169b305 \
+    --hash=sha256:2090f6a85cafc5b2db085124d752757c9d251548cedabe9bd31afe6363e0aff2 \
+    --hash=sha256:20b9b5fbe0b88d0bdef2012ef7dee867f874b72528cf1d08f1d59b0e3850129d \
+    --hash=sha256:22ae2ebf9b0c69d206c003e2f6a914ea33f0a932d4aa16f236afc049d9958f4a \
+    --hash=sha256:22f3105d4fb15c8f57ff3959a58fcab6ce36814486500cd7485651230ad4d4ef \
+    --hash=sha256:23bfd518810af7de1116313ebd9092cb9aa629beb12f6ed631ad53356ed6b86c \
+    --hash=sha256:27e5fc84ccef8dfaabb09d82b7d179c7cf1a3fbc8a966f8274fcb4ab2eb4cadb \
+    --hash=sha256:3380252550e372e8511d49481bd836264c009adb826b23fefcc5dd3c69692f60 \
+    --hash=sha256:3702ea6872c5a2a4eeefa6ffd36b042e9773f05b1f37ae3ef7264b1163c2dcf6 \
+    --hash=sha256:37bb93b2178e02b7b618893990941900fd25b6b9ac0fa49931a40aecdf083fe4 \
+    --hash=sha256:3914f5aaa0f36d5d60e8ece6a308ee1c9784cd75ec8151062614657a114c4478 \
+    --hash=sha256:3a37ffb35399029b45c6cc33640a92bef403c9fd388acce75cdc88f58bd19a81 \
+    --hash=sha256:3c8b88a2ccf5493b6c8da9076fb151ba106960a2df90c2633f342f120751a9e7 \
+    --hash=sha256:3e97b5e938051226dc025ec80980c285b053ffb1e25a3db2a3aa3bc046bf7f56 \
+    --hash=sha256:3ec660d19bbc671e3a6443325f07263be452c453ac9e512f5eb935e7d4ac28b3 \
+    --hash=sha256:3efe2c2cb5763f2f1b275ad2bf7a287d3f7ebbef35648a9726e3b69284a4f3d6 \
+    --hash=sha256:483a6aea59cb89904e1ceabd2b47368b5600fb7de78a6e4a2c2987b2d256cf30 \
+    --hash=sha256:4867cafcbc6585e4b678876c489b9273b13e9fff9f6d6d66add5e15d11d926cb \
+    --hash=sha256:48e171e52d1c4d33888e529b999e5900356b9ae588c2f09a52dcefb158b27506 \
+    --hash=sha256:4a9cb68166a34117d6646c0023c7b759bf197bee5ad4272f420a0141d7eb03a0 \
+    --hash=sha256:4b820514bfc0b98a30e3d85462084779900347e4d49267f747ff54060cc33925 \
+    --hash=sha256:4e18b656c5e844539d506a0a06432274d7bd52a7487e6828c63a63d69185626c \
+    --hash=sha256:4e9f48f58c2c523d5a06faea47866cd35b32655c46b443f163d08c6d0ddb17d6 \
+    --hash=sha256:50b3a2710631848991d0bf7de077502e8994c804bb805aeb2925a981de58ec2e \
+    --hash=sha256:55b6d90641869892caa9ca42ff913f7ff1c5ece06474fbd32fb2cf6834726c95 \
+    --hash=sha256:57feec87371dbb3520da6192213c7d6fc892d5589a93db548331954de8248fd2 \
+    --hash=sha256:58130ecf8f7b8112cdb841486404f1282b9c86ccb30d3519faf301b2e5659133 \
+    --hash=sha256:5845c1fd4866bb5dd3125d89b90e57ed3138241540897de748cdf19de8a2fca2 \
+    --hash=sha256:59bfeae4b25ec05b34f1956eaa1cb38032282cd4dfabc5056d0a1ec4d696d3aa \
+    --hash=sha256:5b48204e8d955c47c55b72779802b219a39acc3ee3d0116d5080c388970b76e3 \
+    --hash=sha256:5c09fcfdccdd0b57867577b719c69e347a436b86cd83747f179dbf0cc0d4c1f3 \
+    --hash=sha256:6180c0ae073bddeb5a97a38c03f30c233e0a4d39cd86166251617d1bbd0af436 \
+    --hash=sha256:682b987361e5fd7a139ed565e30d81fd81e9629acc7d925a205366877d8c8657 \
+    --hash=sha256:6b5d83030255983181005e6cfbac1617ce9746b219bc2aad52201ad121226581 \
+    --hash=sha256:6bb5992037f7a9eff7991ebe4273ea7f51f1c1c511e6a2ce511d0e7bdb754492 \
+    --hash=sha256:73eae06aa53af2ea5270cc066dcaf02cc60d2994bbb2c4ef5764949257d10f43 \
+    --hash=sha256:76f364861c3bfc98cbbcbd402d83454ed9e01a5224bb3a28bf70002a230f73e2 \
+    --hash=sha256:820c661588bd01a0aa62a1283f20d2be4281b086f80dad9e955e690c75fb54a2 \
+    --hash=sha256:82176036e65644a6cc5bd619f65f6f19781e8ec2e5330f51aa9ada7504cc1926 \
+    --hash=sha256:87701f25a2352e5bf7454caa64757642734da9f6b11384c1f9d1a8e699758057 \
+    --hash=sha256:9079dfc6a70abe341f521f78405b8949f96db48da98aeb43f9907f342f627cdc \
+    --hash=sha256:90f8717cb649eea3504091e640a1b8568faad18bd4b9fcd692853a04475a4b80 \
+    --hash=sha256:957cf8e4b6e123a9eea554fa7ebc85674674b713551de587eb318a2df3e00255 \
+    --hash=sha256:99f826cbf970077383d7de805c0681799491cb939c25450b9b5b3ced03ca99f1 \
+    --hash=sha256:9f636b730f7e8cb19feb87094949ba54ee5357440b9658b2a32a5ce4bce53972 \
+    --hash=sha256:a114d03b938376557927ab23f1e950827c3b893ccb94b62fd95d430fd0e5cf53 \
+    --hash=sha256:a185f876e69897a6f3325c3f19f26a297fa058c5e456bfcff8015e9a27e83ae1 \
+    --hash=sha256:a7a9541cd308eed5e30318430a9c74d2132e9a8cb46b901326272d780bf2d423 \
+    --hash=sha256:aa466da5b15ccea564bdab9c89175c762bc12825f4659c11227f515cee76fa4a \
+    --hash=sha256:aaed8b0562be4a0876ee3b6946f6869b7bcdb571a5d1496683505944e268b160 \
+    --hash=sha256:ab7c4ceb38d91570a650dba194e1ca87c2b543488fe9309b4212694174fd539c \
+    --hash=sha256:ac10f4c2b9e770c4e393876e35a7046879d195cd123b4f116d299d442b335bcd \
+    --hash=sha256:b04772ed465fa3cc947db808fa306d79b43e896beb677a56fb2347ca1a49c1fa \
+    --hash=sha256:b1c416351ee6271b2f49b56ad7f308072f6f44b37118d69c2cad94f3fa8a40d5 \
+    --hash=sha256:b225d95519a5bf73860323e633a664b0d85ad3d5bede6d30d95b35d4dfe8805b \
+    --hash=sha256:b2f59caeaf7632cc633b5cf6fc449372b83bbdf0da4ae04d5be36118e46cc0aa \
+    --hash=sha256:b58c621844d55e71c1b7f7c498ce5aa6985d743a1a59034c57a905b3f153c1ef \
+    --hash=sha256:bf6bea52ec97e95560af5ae576bdac3aa3aae0b6758c6efa115236d9e07dae44 \
+    --hash=sha256:c08be4f460903e5a9d0f76818db3250f12e9c344e79314d1d570fc69d7f4eae4 \
+    --hash=sha256:c7053d3b0353a8b9de430a4f4b4268ac9a4fb3481af37dfe49825bf45ca24156 \
+    --hash=sha256:c943a53e9186688b45b323602298ab727d8865d8c9ee0b17f8d62d14b56f0753 \
+    --hash=sha256:ce2186a7df133a9c895dea3331ddc5ddad42cdd0d1ea2f0a51e5d161e4762f28 \
+    --hash=sha256:d093be959277cb7dee84b801eb1af388b6ad3ca6a6b6bf1ed7585895789d027d \
+    --hash=sha256:d094ddec350a2fb899fec68d8353c78233debde9b7d8b4beeafa70825f1c281a \
+    --hash=sha256:d1a9dd711d0877a1ece3d2e4fea11a8e75741ca21954c919406b44e7cf971304 \
+    --hash=sha256:d569388c381b24671589335a3be6e1d45546c2988c2ebe30fdcada8457a31008 \
+    --hash=sha256:d618649d4e70ac6efcbba75be98b26ef5078faad23592f9b51ca492953012429 \
+    --hash=sha256:d83a047959d38a7ff552ff94be767b7fd79b831ad1cd9920662db05fec24fe72 \
+    --hash=sha256:d8fff389528cad1618fb4b26b95550327495462cd745d879a8c7c2115248e399 \
+    --hash=sha256:da1758c76f50c39a2efd5e9859ce7d776317eb1dd34317c8152ac9251fc574a3 \
+    --hash=sha256:db7457bac39421addd0c8449933ac32d8042aae84a14911a757ae6ca3eef1392 \
+    --hash=sha256:e27bbb6d14416713a8bd7aaa1313c0fc8d44ee48d74497a0ff4c3a1b6ccb5167 \
+    --hash=sha256:e617fb6b0b6953fffd762669610c1c4ffd05632c138d61ac7e14ad187870669c \
+    --hash=sha256:e9aa71e15d9d9beaad2c6b9319edcdc0a49a43ef5c0a4c8265ca9ee7d6c67774 \
+    --hash=sha256:ec2abea24d98246b94913b76a125e855eb5c434f7c46546046372fe60f666351 \
+    --hash=sha256:f179dee3b863ab1c59580ff60f9d99f632f34ccb38bf67a33ec6b3ecadd0fd76 \
+    --hash=sha256:f4c035da3f544b1882bac24115f3e2e8760f10a0107614fc9839fd232200b875 \
+    --hash=sha256:f67f217af4b1ff66c68a87318012de788dd95fcfeb24cc889011f4e1c7454dfd \
+    --hash=sha256:f90c822a402cb865e396a504f9fc8173ef34212a342d92e362ca498cad308e28 \
+    --hash=sha256:ff3827aef427c89a25cc96ded1759271a93603aba9fb977a6d264648ebf989db
+    # via
+    #   aiohttp
+    #   yarl
 nest-asyncio==1.6.0 \
     --hash=sha256:6f172d5449aca15afd6c646851f4e31e02c598d553a667e38cafa997cfec55fe \
     --hash=sha256:87af6efd6b5e897c81050477ef65c62e2b2f35d51703cae01aff2905b1852e1c
     # via orbax-checkpoint
+neuralgcm==1.0.0 \
+    --hash=sha256:28f42b73188df92e35d5e187ef9ebd2b2a9922cfc46e05e0ae5202f451ef7c18 \
+    --hash=sha256:e75889894e8948211cec9bf18c89df5e223ef5a6d9a6444bb42f4b36143a5c27
+    # via -r builddeps/test-requirements.txt
+numcodecs==0.13.1 \
+    --hash=sha256:233bc7f26abce24d57e44ea8ebeb5cd17084690b4e7409dd470fdb75528d615f \
+    --hash=sha256:237b7171609e868a20fd313748494444458ccd696062f67e198f7f8f52000c15 \
+    --hash=sha256:2a86f5367af9168e30f99727ff03b27d849c31ad4522060dde0bce2923b3a8bc \
+    --hash=sha256:2eda97dd2f90add98df6d295f2c6ae846043396e3d51a739ca5db6c03b5eb666 \
+    --hash=sha256:3501a848adaddce98a71a262fee15cd3618312692aa419da77acd18af4a6a3f6 \
+    --hash=sha256:3f593c7506b0ab248961a3b13cb148cc6e8355662ff124ac591822310bc55ecf \
+    --hash=sha256:5195bea384a6428f8afcece793860b1ab0ae28143c853f0b2b20d55a8947c917 \
+    --hash=sha256:796b3e6740107e4fa624cc636248a1580138b3f1c579160f260f76ff13a4261b \
+    --hash=sha256:7a60d75179fd6692e301ddfb3b266d51eb598606dcae7b9fc57f986e8d65cb43 \
+    --hash=sha256:80d3071465f03522e776a31045ddf2cfee7f52df468b977ed3afdd7fe5869701 \
+    --hash=sha256:90d3065ae74c9342048ae0046006f99dcb1388b7288da5a19b3bddf9c30c3176 \
+    --hash=sha256:96add4f783c5ce57cc7e650b6cac79dd101daf887c479a00a29bc1487ced180b \
+    --hash=sha256:96e42f73c31b8c24259c5fac6adba0c3ebf95536e37749dc6c62ade2989dca28 \
+    --hash=sha256:a3cf37881df0898f3a9c0d4477df88133fe85185bffe57ba31bcc2fa207709bc \
+    --hash=sha256:da2230484e6102e5fa3cc1a5dd37ca1f92dfbd183d91662074d6f7574e3e8f53 \
+    --hash=sha256:e5db4824ebd5389ea30e54bc8aeccb82d514d28b6b68da6c536b8fa4596f4bca \
+    --hash=sha256:eda7d7823c9282e65234731fd6bd3986b1f9e035755f7fed248d7d366bb291ab
+    # via zarr
 numpy==2.1.0 \
     --hash=sha256:08801848a40aea24ce16c2ecde3b756f9ad756586fb2d13210939eb69b023f5b \
     --hash=sha256:0937e54c09f7a9a68da6889362ddd2ff584c02d015ec92672c099b61555f8911 \
@@ -425,6 +919,7 @@ numpy==2.1.0 \
     # via
     #   -r builddeps/test-requirements.txt
     #   chex
+    #   dinosaur-dycore
     #   dm-haiku
     #   e3nn-jax
     #   flax
@@ -434,11 +929,19 @@ numpy==2.1.0 \
     #   jmp
     #   jraph
     #   ml-dtypes
+    #   neuralgcm
+    #   numcodecs
     #   opt-einsum
     #   optax
     #   orbax-checkpoint
+    #   pandas
+    #   scikit-learn
     #   scipy
+    #   tensorflow-probability
     #   tensorstore
+    #   xarray
+    #   xarray-tensorstore
+    #   zarr
 nvidia-cublas-cu12==12.6.1.4 \
     --hash=sha256:5dd125ece5469dbdceebe2e9536ad8fc4abd38aa394a7ace42fc8a930a1e81e3 \
     --hash=sha256:5e5d384583d72ac364064ced3dd92a5caa59a8a57568595c9f82e83d255b2481 \
@@ -496,6 +999,10 @@ nvidia-nvjitlink-cu12==12.6.68 \
     #   nvidia-cufft-cu12
     #   nvidia-cusolver-cu12
     #   nvidia-cusparse-cu12
+oauthlib==3.2.2 \
+    --hash=sha256:8139f29aac13e25d502680e9e19963e83f16838d48a0d71c287fe40e7067fbca \
+    --hash=sha256:9859c40929662bec5d64f34d01c99e093149682a3f38915dc0655d5a633dd918
+    # via requests-oauthlib
 opt-einsum==3.3.0 \
     --hash=sha256:2455e59e3947d3c275477df7f5205b30635e266fe6dc300e3d9f9646bfcea147 \
     --hash=sha256:59f6475f77bbc37dcf7cd748519c0ec60722e91e63ca114e68821c0c54a46549
@@ -506,10 +1013,174 @@ optax==0.2.3 \
     # via
     #   flax
     #   jax-md
+    #   neuralgcm
 orbax-checkpoint==0.6.4 \
     --hash=sha256:366b4d528a7322e1b3d9ddcaed45c8515add0d2fc69c8975c30d98638543240f \
     --hash=sha256:b4f2608ee4d1da67f7619fc35ff9c928ecdf4ccf7546eeb43ecf38c2608b6dea
     # via flax
+packaging==24.1 \
+    --hash=sha256:026ed72c8ed3fcce5bf8950572258698927fd1dbda10a5e981cdf0ac37f4f002 \
+    --hash=sha256:5b8f2217dbdbd2f7f384c41c628544e6d52f2d0f53c6d0c3ea61aa5d1d7ff124
+    # via xarray
+pandas==2.2.3 \
+    --hash=sha256:062309c1b9ea12a50e8ce661145c6aab431b1e99530d3cd60640e255778bd43a \
+    --hash=sha256:15c0e1e02e93116177d29ff83e8b1619c93ddc9c49083f237d4312337a61165d \
+    --hash=sha256:1948ddde24197a0f7add2bdc4ca83bf2b1ef84a1bc8ccffd95eda17fd836ecb5 \
+    --hash=sha256:1db71525a1538b30142094edb9adc10be3f3e176748cd7acc2240c2f2e5aa3a4 \
+    --hash=sha256:22a9d949bfc9a502d320aa04e5d02feab689d61da4e7764b62c30b991c42c5f0 \
+    --hash=sha256:29401dbfa9ad77319367d36940cd8a0b3a11aba16063e39632d98b0e931ddf32 \
+    --hash=sha256:31d0ced62d4ea3e231a9f228366919a5ea0b07440d9d4dac345376fd8e1477ea \
+    --hash=sha256:3508d914817e153ad359d7e069d752cdd736a247c322d932eb89e6bc84217f28 \
+    --hash=sha256:37e0aced3e8f539eccf2e099f65cdb9c8aa85109b0be6e93e2baff94264bdc6f \
+    --hash=sha256:381175499d3802cde0eabbaf6324cce0c4f5d52ca6f8c377c29ad442f50f6348 \
+    --hash=sha256:38cf8125c40dae9d5acc10fa66af8ea6fdf760b2714ee482ca691fc66e6fcb18 \
+    --hash=sha256:3b71f27954685ee685317063bf13c7709a7ba74fc996b84fc6821c59b0f06468 \
+    --hash=sha256:3fc6873a41186404dad67245896a6e440baacc92f5b716ccd1bc9ed2995ab2c5 \
+    --hash=sha256:4850ba03528b6dd51d6c5d273c46f183f39a9baf3f0143e566b89450965b105e \
+    --hash=sha256:4f18ba62b61d7e192368b84517265a99b4d7ee8912f8708660fb4a366cc82667 \
+    --hash=sha256:56534ce0746a58afaf7942ba4863e0ef81c9c50d3f0ae93e9497d6a41a057645 \
+    --hash=sha256:59ef3764d0fe818125a5097d2ae867ca3fa64df032331b7e0917cf5d7bf66b13 \
+    --hash=sha256:5dbca4c1acd72e8eeef4753eeca07de9b1db4f398669d5994086f788a5d7cc30 \
+    --hash=sha256:5de54125a92bb4d1c051c0659e6fcb75256bf799a732a87184e5ea503965bce3 \
+    --hash=sha256:61c5ad4043f791b61dd4752191d9f07f0ae412515d59ba8f005832a532f8736d \
+    --hash=sha256:6374c452ff3ec675a8f46fd9ab25c4ad0ba590b71cf0656f8b6daa5202bca3fb \
+    --hash=sha256:63cc132e40a2e084cf01adf0775b15ac515ba905d7dcca47e9a251819c575ef3 \
+    --hash=sha256:66108071e1b935240e74525006034333f98bcdb87ea116de573a6a0dccb6c039 \
+    --hash=sha256:6dfcb5ee8d4d50c06a51c2fffa6cff6272098ad6540aed1a76d15fb9318194d8 \
+    --hash=sha256:7c2875855b0ff77b2a64a0365e24455d9990730d6431b9e0ee18ad8acee13dbd \
+    --hash=sha256:7eee9e7cea6adf3e3d24e304ac6b8300646e2a5d1cd3a3c2abed9101b0846761 \
+    --hash=sha256:800250ecdadb6d9c78eae4990da62743b857b470883fa27f652db8bdde7f6659 \
+    --hash=sha256:86976a1c5b25ae3f8ccae3a5306e443569ee3c3faf444dfd0f41cda24667ad57 \
+    --hash=sha256:8cd6d7cc958a3910f934ea8dbdf17b2364827bb4dafc38ce6eef6bb3d65ff09c \
+    --hash=sha256:99df71520d25fade9db7c1076ac94eb994f4d2673ef2aa2e86ee039b6746d20c \
+    --hash=sha256:a5a1595fe639f5988ba6a8e5bc9649af3baf26df3998a0abe56c02609392e0a4 \
+    --hash=sha256:ad5b65698ab28ed8d7f18790a0dc58005c7629f227be9ecc1072aa74c0c1d43a \
+    --hash=sha256:b1d432e8d08679a40e2a6d8b2f9770a5c21793a6f9f47fdd52c5ce1948a5a8a9 \
+    --hash=sha256:b8661b0238a69d7aafe156b7fa86c44b881387509653fdf857bebc5e4008ad42 \
+    --hash=sha256:ba96630bc17c875161df3818780af30e43be9b166ce51c9a18c1feae342906c2 \
+    --hash=sha256:bc6b93f9b966093cb0fd62ff1a7e4c09e6d546ad7c1de191767baffc57628f39 \
+    --hash=sha256:c124333816c3a9b03fbeef3a9f230ba9a737e9e5bb4060aa2107a86cc0a497fc \
+    --hash=sha256:cd8d0c3be0515c12fed0bdbae072551c8b54b7192c7b1fda0ba56059a0179698 \
+    --hash=sha256:d9c45366def9a3dd85a6454c0e7908f2b3b8e9c138f5dc38fed7ce720d8453ed \
+    --hash=sha256:f00d1345d84d8c86a63e476bb4955e46458b304b9575dcf71102b5c705320015 \
+    --hash=sha256:f3a255b2c19987fbbe62a9dfd6cff7ff2aa9ccab3fc75218fd4b7530f01efa24 \
+    --hash=sha256:fffb8ae78d8af97f849404f21411c95062db1496aeb3e56f146f0355c9989319
+    # via
+    #   dinosaur-dycore
+    #   neuralgcm
+    #   xarray
+pint==0.24.4 \
+    --hash=sha256:35275439b574837a6cd3020a5a4a73645eb125ce4152a73a2f126bf164b91b80 \
+    --hash=sha256:aa54926c8772159fcf65f82cc0d34de6768c151b32ad1deb0331291c38fe7659
+    # via dinosaur-dycore
+platformdirs==4.3.6 \
+    --hash=sha256:357fb2acbc885b0419afd3ce3ed34564c13c9b95c89360cd9563f73aa5e2b907 \
+    --hash=sha256:73e575e1408ab8103900836b97580d5307456908a03e92031bab39e4554cc3fb
+    # via pint
+propcache==0.2.0 \
+    --hash=sha256:00181262b17e517df2cd85656fcd6b4e70946fe62cd625b9d74ac9977b64d8d9 \
+    --hash=sha256:0e53cb83fdd61cbd67202735e6a6687a7b491c8742dfc39c9e01e80354956763 \
+    --hash=sha256:1235c01ddaa80da8235741e80815ce381c5267f96cc49b1477fdcf8c047ef325 \
+    --hash=sha256:140fbf08ab3588b3468932974a9331aff43c0ab8a2ec2c608b6d7d1756dbb6cb \
+    --hash=sha256:191db28dc6dcd29d1a3e063c3be0b40688ed76434622c53a284e5427565bbd9b \
+    --hash=sha256:1e41d67757ff4fbc8ef2af99b338bfb955010444b92929e9e55a6d4dcc3c4f09 \
+    --hash=sha256:1ec43d76b9677637a89d6ab86e1fef70d739217fefa208c65352ecf0282be957 \
+    --hash=sha256:20a617c776f520c3875cf4511e0d1db847a076d720714ae35ffe0df3e440be68 \
+    --hash=sha256:218db2a3c297a3768c11a34812e63b3ac1c3234c3a086def9c0fee50d35add1f \
+    --hash=sha256:22aa8f2272d81d9317ff5756bb108021a056805ce63dd3630e27d042c8092798 \
+    --hash=sha256:25a1f88b471b3bc911d18b935ecb7115dff3a192b6fef46f0bfaf71ff4f12418 \
+    --hash=sha256:25c8d773a62ce0451b020c7b29a35cfbc05de8b291163a7a0f3b7904f27253e6 \
+    --hash=sha256:2a60ad3e2553a74168d275a0ef35e8c0a965448ffbc3b300ab3a5bb9956c2162 \
+    --hash=sha256:2a66df3d4992bc1d725b9aa803e8c5a66c010c65c741ad901e260ece77f58d2f \
+    --hash=sha256:2ccc28197af5313706511fab3a8b66dcd6da067a1331372c82ea1cb74285e036 \
+    --hash=sha256:2e900bad2a8456d00a113cad8c13343f3b1f327534e3589acc2219729237a2e8 \
+    --hash=sha256:2ee7606193fb267be4b2e3b32714f2d58cad27217638db98a60f9efb5efeccc2 \
+    --hash=sha256:33ac8f098df0585c0b53009f039dfd913b38c1d2edafed0cedcc0c32a05aa110 \
+    --hash=sha256:3444cdba6628accf384e349014084b1cacd866fbb88433cd9d279d90a54e0b23 \
+    --hash=sha256:363ea8cd3c5cb6679f1c2f5f1f9669587361c062e4899fce56758efa928728f8 \
+    --hash=sha256:375a12d7556d462dc64d70475a9ee5982465fbb3d2b364f16b86ba9135793638 \
+    --hash=sha256:388f3217649d6d59292b722d940d4d2e1e6a7003259eb835724092a1cca0203a \
+    --hash=sha256:3947483a381259c06921612550867b37d22e1df6d6d7e8361264b6d037595f44 \
+    --hash=sha256:39e104da444a34830751715f45ef9fc537475ba21b7f1f5b0f4d71a3b60d7fe2 \
+    --hash=sha256:3c997f8c44ec9b9b0bcbf2d422cc00a1d9b9c681f56efa6ca149a941e5560da2 \
+    --hash=sha256:3dfafb44f7bb35c0c06eda6b2ab4bfd58f02729e7c4045e179f9a861b07c9850 \
+    --hash=sha256:3ebbcf2a07621f29638799828b8d8668c421bfb94c6cb04269130d8de4fb7136 \
+    --hash=sha256:3f88a4095e913f98988f5b338c1d4d5d07dbb0b6bad19892fd447484e483ba6b \
+    --hash=sha256:439e76255daa0f8151d3cb325f6dd4a3e93043e6403e6491813bcaaaa8733887 \
+    --hash=sha256:4569158070180c3855e9c0791c56be3ceeb192defa2cdf6a3f39e54319e56b89 \
+    --hash=sha256:466c219deee4536fbc83c08d09115249db301550625c7fef1c5563a584c9bc87 \
+    --hash=sha256:4a9d9b4d0a9b38d1c391bb4ad24aa65f306c6f01b512e10a8a34a2dc5675d348 \
+    --hash=sha256:4c7dde9e533c0a49d802b4f3f218fa9ad0a1ce21f2c2eb80d5216565202acab4 \
+    --hash=sha256:53d1bd3f979ed529f0805dd35ddaca330f80a9a6d90bc0121d2ff398f8ed8861 \
+    --hash=sha256:55346705687dbd7ef0d77883ab4f6fabc48232f587925bdaf95219bae072491e \
+    --hash=sha256:56295eb1e5f3aecd516d91b00cfd8bf3a13991de5a479df9e27dd569ea23959c \
+    --hash=sha256:56bb5c98f058a41bb58eead194b4db8c05b088c93d94d5161728515bd52b052b \
+    --hash=sha256:5a5b3bb545ead161be780ee85a2b54fdf7092815995661947812dde94a40f6fb \
+    --hash=sha256:5f2564ec89058ee7c7989a7b719115bdfe2a2fb8e7a4543b8d1c0cc4cf6478c1 \
+    --hash=sha256:608cce1da6f2672a56b24a015b42db4ac612ee709f3d29f27a00c943d9e851de \
+    --hash=sha256:63f13bf09cc3336eb04a837490b8f332e0db41da66995c9fd1ba04552e516354 \
+    --hash=sha256:662dd62358bdeaca0aee5761de8727cfd6861432e3bb828dc2a693aa0471a563 \
+    --hash=sha256:676135dcf3262c9c5081cc8f19ad55c8a64e3f7282a21266d05544450bffc3a5 \
+    --hash=sha256:67aeb72e0f482709991aa91345a831d0b707d16b0257e8ef88a2ad246a7280bf \
+    --hash=sha256:67b69535c870670c9f9b14a75d28baa32221d06f6b6fa6f77a0a13c5a7b0a5b9 \
+    --hash=sha256:682a7c79a2fbf40f5dbb1eb6bfe2cd865376deeac65acf9beb607505dced9e12 \
+    --hash=sha256:6994984550eaf25dd7fc7bd1b700ff45c894149341725bb4edc67f0ffa94efa4 \
+    --hash=sha256:69d3a98eebae99a420d4b28756c8ce6ea5a29291baf2dc9ff9414b42676f61d5 \
+    --hash=sha256:6e2e54267980349b723cff366d1e29b138b9a60fa376664a157a342689553f71 \
+    --hash=sha256:73e4b40ea0eda421b115248d7e79b59214411109a5bc47d0d48e4c73e3b8fcf9 \
+    --hash=sha256:74acd6e291f885678631b7ebc85d2d4aec458dd849b8c841b57ef04047833bed \
+    --hash=sha256:7665f04d0c7f26ff8bb534e1c65068409bf4687aa2534faf7104d7182debb336 \
+    --hash=sha256:7735e82e3498c27bcb2d17cb65d62c14f1100b71723b68362872bca7d0913d90 \
+    --hash=sha256:77a86c261679ea5f3896ec060be9dc8e365788248cc1e049632a1be682442063 \
+    --hash=sha256:7cf18abf9764746b9c8704774d8b06714bcb0a63641518a3a89c7f85cc02c2ad \
+    --hash=sha256:83928404adf8fb3d26793665633ea79b7361efa0287dfbd372a7e74311d51ee6 \
+    --hash=sha256:8e40876731f99b6f3c897b66b803c9e1c07a989b366c6b5b475fafd1f7ba3fb8 \
+    --hash=sha256:8f188cfcc64fb1266f4684206c9de0e80f54622c3f22a910cbd200478aeae61e \
+    --hash=sha256:91997d9cb4a325b60d4e3f20967f8eb08dfcb32b22554d5ef78e6fd1dda743a2 \
+    --hash=sha256:91ee8fc02ca52e24bcb77b234f22afc03288e1dafbb1f88fe24db308910c4ac7 \
+    --hash=sha256:92fe151145a990c22cbccf9ae15cae8ae9eddabfc949a219c9f667877e40853d \
+    --hash=sha256:945db8ee295d3af9dbdbb698cce9bbc5c59b5c3fe328bbc4387f59a8a35f998d \
+    --hash=sha256:9517d5e9e0731957468c29dbfd0f976736a0e55afaea843726e887f36fe017df \
+    --hash=sha256:952e0d9d07609d9c5be361f33b0d6d650cd2bae393aabb11d9b719364521984b \
+    --hash=sha256:97a58a28bcf63284e8b4d7b460cbee1edaab24634e82059c7b8c09e65284f178 \
+    --hash=sha256:97e48e8875e6c13909c800fa344cd54cc4b2b0db1d5f911f840458a500fde2c2 \
+    --hash=sha256:9e0f07b42d2a50c7dd2d8675d50f7343d998c64008f1da5fef888396b7f84630 \
+    --hash=sha256:a3dc1a4b165283bd865e8f8cb5f0c64c05001e0718ed06250d8cac9bec115b48 \
+    --hash=sha256:a3ebe9a75be7ab0b7da2464a77bb27febcb4fab46a34f9288f39d74833db7f61 \
+    --hash=sha256:a64e32f8bd94c105cc27f42d3b658902b5bcc947ece3c8fe7bc1b05982f60e89 \
+    --hash=sha256:a6ed8db0a556343d566a5c124ee483ae113acc9a557a807d439bcecc44e7dfbb \
+    --hash=sha256:ad9c9b99b05f163109466638bd30ada1722abb01bbb85c739c50b6dc11f92dc3 \
+    --hash=sha256:b33d7a286c0dc1a15f5fc864cc48ae92a846df287ceac2dd499926c3801054a6 \
+    --hash=sha256:bc092ba439d91df90aea38168e11f75c655880c12782facf5cf9c00f3d42b562 \
+    --hash=sha256:c436130cc779806bdf5d5fae0d848713105472b8566b75ff70048c47d3961c5b \
+    --hash=sha256:c5869b8fd70b81835a6f187c5fdbe67917a04d7e52b6e7cc4e5fe39d55c39d58 \
+    --hash=sha256:c5ecca8f9bab618340c8e848d340baf68bcd8ad90a8ecd7a4524a81c1764b3db \
+    --hash=sha256:cfac69017ef97db2438efb854edf24f5a29fd09a536ff3a992b75990720cdc99 \
+    --hash=sha256:d2f0d0f976985f85dfb5f3d685697ef769faa6b71993b46b295cdbbd6be8cc37 \
+    --hash=sha256:d5bed7f9805cc29c780f3aee05de3262ee7ce1f47083cfe9f77471e9d6777e83 \
+    --hash=sha256:d6a21ef516d36909931a2967621eecb256018aeb11fc48656e3257e73e2e247a \
+    --hash=sha256:d9b6ddac6408194e934002a69bcaadbc88c10b5f38fb9307779d1c629181815d \
+    --hash=sha256:db47514ffdbd91ccdc7e6f8407aac4ee94cc871b15b577c1c324236b013ddd04 \
+    --hash=sha256:df81779732feb9d01e5d513fad0122efb3d53bbc75f61b2a4f29a020bc985e70 \
+    --hash=sha256:e4a91d44379f45f5e540971d41e4626dacd7f01004826a18cb048e7da7e96544 \
+    --hash=sha256:e63e3e1e0271f374ed489ff5ee73d4b6e7c60710e1f76af5f0e1a6117cd26394 \
+    --hash=sha256:e70fac33e8b4ac63dfc4c956fd7d85a0b1139adcfc0d964ce288b7c527537fea \
+    --hash=sha256:ecddc221a077a8132cf7c747d5352a15ed763b674c0448d811f408bf803d9ad7 \
+    --hash=sha256:f45eec587dafd4b2d41ac189c2156461ebd0c1082d2fe7013571598abb8505d1 \
+    --hash=sha256:f52a68c21363c45297aca15561812d542f8fc683c85201df0bebe209e349f793 \
+    --hash=sha256:f571aea50ba5623c308aa146eb650eebf7dbe0fd8c5d946e28343cb3b5aad577 \
+    --hash=sha256:f60f0ac7005b9f5a6091009b09a419ace1610e163fa5deaba5ce3484341840e7 \
+    --hash=sha256:f6475a1b2ecb310c98c28d271a30df74f9dd436ee46d09236a6b750a7599ce57 \
+    --hash=sha256:f6d5749fdd33d90e34c2efb174c7e236829147a2713334d708746e94c4bde40d \
+    --hash=sha256:f902804113e032e2cdf8c71015651c97af6418363bea8d78dc0911d56c335032 \
+    --hash=sha256:fa1076244f54bb76e65e22cb6910365779d5c3d71d1f18b275f1dfc7b0d71b4d \
+    --hash=sha256:fc2db02409338bf36590aa985a461b2c96fce91f8e7e0f14c50c5fcc4f229016 \
+    --hash=sha256:ffcad6c564fe6b9b8916c1aefbb37a362deebf9394bd2974e9d84232e3e08504
+    # via yarl
+proto-plus==1.25.0 \
+    --hash=sha256:c91fc4a65074ade8e458e95ef8bac34d4008daa7cce4a12d6707066fca648961 \
+    --hash=sha256:fbb17f57f7bd05a68b7707e745e26528b0b3c34e378db91eef93912c54982d91
+    # via google-api-core
 protobuf==5.28.2 \
     --hash=sha256:2c69461a7fcc8e24be697624c09a839976d82ae75062b11a0972e41fd2cd9132 \
     --hash=sha256:35cfcb15f213449af7ff6198d6eb5f739c37d7e4f1c09b5d0641babf2cc0c68f \
@@ -522,11 +1193,33 @@ protobuf==5.28.2 \
     --hash=sha256:c0ea0123dac3399a2eeb1a1443d82b7afc9ff40241433296769f7da42d142ec3 \
     --hash=sha256:ca53faf29896c526863366a52a8f4d88e69cd04ec9571ed6082fa117fac3ab36 \
     --hash=sha256:eeea10f3dc0ac7e6b4933d32db20662902b4ab81bf28df12218aa389e9c2102d
-    # via orbax-checkpoint
+    # via
+    #   google-api-core
+    #   googleapis-common-protos
+    #   orbax-checkpoint
+    #   proto-plus
+pyasn1==0.6.1 \
+    --hash=sha256:0d632f46f2ba09143da3a8afe9e33fb6f92fa2320ab7e886e2d0f7672af84629 \
+    --hash=sha256:6f580d2bdd84365380830acf45550f2511469f673cb4a5ae3857a3170128b034
+    # via
+    #   pyasn1-modules
+    #   rsa
+pyasn1-modules==0.4.1 \
+    --hash=sha256:49bfa96b45a292b711e986f222502c1c9a5e1f4e568fc30e2574a6c7d07838fd \
+    --hash=sha256:c28e2dbf9c06ad61c71a075c7e0f9fd0f1b0bb2d2ad4377f240d33ac2ab60a7c
+    # via google-auth
 pygments==2.18.0 \
     --hash=sha256:786ff802f32e91311bff3889f6e9a86e81505fe99f2735bb6d60ae0c5004f199 \
     --hash=sha256:b8e6aca0523f3ab76fee51799c488e38782ac06eafcf95e7ba832985c8e7b13a
     # via rich
+python-dateutil==2.9.0.post0 \
+    --hash=sha256:37dd54208da7e1cd875388217d5e00ebd4179249f90fb72437e91a35459a0ad3 \
+    --hash=sha256:a8b2bc7bffae282281c8140a97d3aa9c14da0b136dfe83f850eea9a5f7470427
+    # via pandas
+pytz==2024.2 \
+    --hash=sha256:2aa355083c50a0f93fa581709deac0c9ad65cca8a9e9beac660adcbd493c798a \
+    --hash=sha256:31c7c1817eb7fae7ca4b8c7ee50c72f93aa2dd863de768e1ef4245d426aa0725
+    # via pandas
 pyyaml==6.0.2 \
     --hash=sha256:01179a4a8559ab5de078078f37e5c1a30d76bb88519906844fd7bdea1b7729ff \
     --hash=sha256:0833f8694549e586547b576dcfaba4a6b55b9e96098b36cdc7ebefe667dfed48 \
@@ -588,11 +1281,52 @@ pyyaml==6.0.2 \
 requests==2.32.3 ; sys_platform == "linux" \
     --hash=sha256:55365417734eb18255590a9ff9eb97e9e1da868d4ccd6402399eaf68af20a760 \
     --hash=sha256:70761cfe03c773ceb22aa2f671b4757976145175cdfca038c02654d061d6dcc6
-    # via -r builddeps/test-requirements.txt
+    # via
+    #   -r builddeps/test-requirements.txt
+    #   gcsfs
+    #   google-api-core
+    #   google-cloud-storage
+    #   requests-oauthlib
+requests-oauthlib==2.0.0 \
+    --hash=sha256:7dd8a5c40426b779b0868c404bdef9768deccf22749cde15852df527e6269b36 \
+    --hash=sha256:b3dffaebd884d8cd778494369603a9e7b58d29111bf6b41bdc2dcd87203af4e9
+    # via google-auth-oauthlib
 rich==13.9.2 \
     --hash=sha256:51a2c62057461aaf7152b4d611168f93a9fc73068f8ded2790f29fe2b5366d0c \
     --hash=sha256:8c82a3d3f8dcfe9e734771313e606b39d8247bb6b826e196f4914b333b743cf1
     # via flax
+rsa==4.9 \
+    --hash=sha256:90260d9058e514786967344d0ef75fa8727eed8a7d2e43ce9f4bcf1b536174f7 \
+    --hash=sha256:e38464a49c6c85d7f1351b0126661487a7e0a14a50f1675ec50eb34d4f20ef21
+    # via google-auth
+scikit-learn==1.5.2 \
+    --hash=sha256:03b6158efa3faaf1feea3faa884c840ebd61b6484167c711548fce208ea09445 \
+    --hash=sha256:178ddd0a5cb0044464fc1bfc4cca5b1833bfc7bb022d70b05db8530da4bb3dd3 \
+    --hash=sha256:1ff45e26928d3b4eb767a8f14a9a6efbf1cbff7c05d1fb0f95f211a89fd4f5de \
+    --hash=sha256:299406827fb9a4f862626d0fe6c122f5f87f8910b86fe5daa4c32dcd742139b6 \
+    --hash=sha256:2d4cad1119c77930b235579ad0dc25e65c917e756fe80cab96aa3b9428bd3fb0 \
+    --hash=sha256:394397841449853c2290a32050382edaec3da89e35b3e03d6cc966aebc6a8ae6 \
+    --hash=sha256:3a686885a4b3818d9e62904d91b57fa757fc2bed3e465c8b177be652f4dd37c8 \
+    --hash=sha256:3b923d119d65b7bd555c73be5423bf06c0105678ce7e1f558cb4b40b0a5502b1 \
+    --hash=sha256:3bed4909ba187aca80580fe2ef370d9180dcf18e621a27c4cf2ef10d279a7efe \
+    --hash=sha256:52788f48b5d8bca5c0736c175fa6bdaab2ef00a8f536cda698db61bd89c551c1 \
+    --hash=sha256:57cc1786cfd6bd118220a92ede80270132aa353647684efa385a74244a41e3b1 \
+    --hash=sha256:643964678f4b5fbdc95cbf8aec638acc7aa70f5f79ee2cdad1eec3df4ba6ead8 \
+    --hash=sha256:6c16d84a0d45e4894832b3c4d0bf73050939e21b99b01b6fd59cbb0cf39163b6 \
+    --hash=sha256:757c7d514ddb00ae249832fe87100d9c73c6ea91423802872d9e74970a0e40b9 \
+    --hash=sha256:8c412ccc2ad9bf3755915e3908e677b367ebc8d010acbb3f182814524f2e5540 \
+    --hash=sha256:b0768ad641981f5d3a198430a1d31c3e044ed2e8a6f22166b4d546a5116d7908 \
+    --hash=sha256:b4237ed7b3fdd0a4882792e68ef2545d5baa50aca3bb45aa7df468138ad8f94d \
+    --hash=sha256:b7b0f9a0b1040830d38c39b91b3a44e1b643f4b36e36567b80b7c6bd2202a27f \
+    --hash=sha256:c15b1ca23d7c5f33cc2cb0a0d6aaacf893792271cddff0edbd6a40e8319bc113 \
+    --hash=sha256:ca64b3089a6d9b9363cd3546f8978229dcbb737aceb2c12144ee3f70f95684b7 \
+    --hash=sha256:e9a702e2de732bbb20d3bad29ebd77fc05a6b427dc49964300340e4c9328b3f5 \
+    --hash=sha256:f60021ec1574e56632be2a36b946f8143bf4e5e6af4a06d85281adc22938e0dd \
+    --hash=sha256:f7284ade780084d94505632241bf78c44ab3b6f1e8ccab3d2af58e0e950f9c12 \
+    --hash=sha256:f763897fe92d0e903aa4847b0aec0e68cadfff77e8a0687cabd946c89d17e675 \
+    --hash=sha256:f8b0ccd4a902836493e026c03256e8b206656f91fbcc4fde28c57a5b752561f1 \
+    --hash=sha256:f932a02c3f4956dfb981391ab24bda1dbd90fe3d628e4b42caef3e041c67707a
+    # via dinosaur-dycore
 scipy==1.14.1 \
     --hash=sha256:0c2f95de3b04e26f5f3ad5bb05e74ba7f68b837133a4492414b3afd79dfe540e \
     --hash=sha256:1729560c906963fc8389f6aac023739ff3983e727b1a4d87696b7bf108316a79 \
@@ -628,12 +1362,17 @@ scipy==1.14.1 \
     --hash=sha256:edaf02b82cd7639db00dbff629995ef185c8df4c3ffa71a5562a595765a06ce1 \
     --hash=sha256:fef8c87f8abfb884dac04e97824b61299880c43f4ce675dd2cbeadd3c9b466d2
     # via
+    #   dinosaur-dycore
     #   jax
     #   jaxlib
+    #   scikit-learn
 six==1.16.0 \
     --hash=sha256:1e61c37477a1626458e36f7b1d82aa5c9b094fa4802892072e49de9c60c4c926 \
     --hash=sha256:8abb2f1d86890a2dfb989f9a77cfcfd3e47c2a354b01111771326f8aa26e0254
-    # via ml-collections
+    # via
+    #   ml-collections
+    #   python-dateutil
+    #   tensorflow-probability
 sympy==1.13.3 \
     --hash=sha256:54612cf55a62755ee71824ce692986f23c88ffa77207b30c1368eda4a7060f73 \
     --hash=sha256:b27fd2c6530e0ab39e275fc9b683895367e51d5da91baa8d3d64db2565fec4d9
@@ -642,6 +1381,9 @@ tabulate==0.9.0 \
     --hash=sha256:0095b12bf5966de529c0feb1fa08671671b3368eec77d7ef7ab114be2c068b3c \
     --hash=sha256:024ca478df22e9340661486f85298cff5f6dcdba14f3813e8830015b9ed1948f
     # via dm-haiku
+tensorflow-probability==0.24.0 \
+    --hash=sha256:8c1774683e38359dbcaf3697e79b7e6a4e69b9c7b3679e78ee18f43e59e5759b
+    # via neuralgcm
 tensorstore==0.1.66 \
     --hash=sha256:0261b1b7121186ac4898c9b4e29289c46754f9d7d73e906620e9a21386fe564e \
     --hash=sha256:06f7dba630d75f65201b1c021fcd89e5f8538506b444b5dd712f35c11b812fc8 \
@@ -667,10 +1409,19 @@ tensorstore==0.1.66 \
     # via
     #   flax
     #   orbax-checkpoint
+    #   xarray-tensorstore
+threadpoolctl==3.5.0 \
+    --hash=sha256:082433502dd922bf738de0d8bcc4fdcbf0979ff44c42bd40f5af8a282f6fa107 \
+    --hash=sha256:56c1e26c150397e58c4926da8eeee87533b1e32bef131bd4bf6a2f45f3185467
+    # via scikit-learn
 toolz==1.0.0 \
     --hash=sha256:292c8f1c4e7516bf9086f8850935c799a874039c8bcf959d47b600e4c44a6236 \
     --hash=sha256:2c86e3d9a04798ac556793bced838816296a2f085017664e4995cb40a1047a02
     # via chex
+tree-math==0.2.1 \
+    --hash=sha256:77613c2360e067b16c0738b32e7d64aa5e1167e609ca1a32b6fa05dcb0d330bb \
+    --hash=sha256:80b01a5d29c1339db0cd8b465f9fad4f872adac117ed599b14f093a6c52821fe
+    # via dinosaur-dycore
 typing-extensions==4.12.2 \
     --hash=sha256:04e5ca0351e0f3f85c6853954072df659d0d13fac324d0072316b67d7794700d \
     --hash=sha256:1a7ead55c7e559dd4dee8856e3a88b41225abfe1ce8df57b7c13915fe121ffb8
@@ -678,11 +1429,117 @@ typing-extensions==4.12.2 \
     #   chex
     #   etils
     #   flax
+    #   flexcache
+    #   flexparser
     #   orbax-checkpoint
+    #   pint
+tzdata==2024.2 \
+    --hash=sha256:7d85cc416e9382e69095b7bdf4afd9e3880418a2413feec7069d533d6b4e31cc \
+    --hash=sha256:a48093786cdcde33cad18c2555e8532f34422074448fbc874186f0abd79565cd
+    # via pandas
 urllib3==2.2.2 \
     --hash=sha256:a448b2f64d686155468037e1ace9f2d2199776e17f0a46610480d311f73e3472 \
     --hash=sha256:dd505485549a7a552833da5e6063639d0d177c04f23bc3864e41e5dc5f612168
     # via requests
+xarray==2024.10.0 \
+    --hash=sha256:ae1d38cb44a0324dfb61e492394158ae22389bf7de9f3c174309c17376df63a0 \
+    --hash=sha256:e369e2bac430e418c2448e5b96f07da4635f98c1319aa23cfeb3fbcb9a01d2e0
+    # via
+    #   dinosaur-dycore
+    #   neuralgcm
+    #   xarray-tensorstore
+xarray-tensorstore==0.1.4 \
+    --hash=sha256:1dacec6376912373dbf9b239d96184090f400a49cd66017f5bf0c824793c9552 \
+    --hash=sha256:de3558c08891c4a7524cf579cb271d6cc005962efd1b8601dd679b58b13a68c8
+    # via dinosaur-dycore
+yarl==1.17.1 \
+    --hash=sha256:06157fb3c58f2736a5e47c8fcbe1afc8b5de6fb28b14d25574af9e62150fcaac \
+    --hash=sha256:067a63fcfda82da6b198fa73079b1ca40b7c9b7994995b6ee38acda728b64d47 \
+    --hash=sha256:0b1794853124e2f663f0ea54efb0340b457f08d40a1cef78edfa086576179c91 \
+    --hash=sha256:0bdff5e0995522706c53078f531fb586f56de9c4c81c243865dd5c66c132c3b5 \
+    --hash=sha256:117ed8b3732528a1e41af3aa6d4e08483c2f0f2e3d3d7dca7cf538b3516d93df \
+    --hash=sha256:14bc88baa44e1f84164a392827b5defb4fa8e56b93fecac3d15315e7c8e5d8b3 \
+    --hash=sha256:1654ec814b18be1af2c857aa9000de7a601400bd4c9ca24629b18486c2e35463 \
+    --hash=sha256:16bca6678a83657dd48df84b51bd56a6c6bd401853aef6d09dc2506a78484c7b \
+    --hash=sha256:1a3b91c44efa29e6c8ef8a9a2b583347998e2ba52c5d8280dbd5919c02dfc3b5 \
+    --hash=sha256:1a52a1ffdd824fb1835272e125385c32fd8b17fbdefeedcb4d543cc23b332d74 \
+    --hash=sha256:1ce36ded585f45b1e9bb36d0ae94765c6608b43bd2e7f5f88079f7a85c61a4d3 \
+    --hash=sha256:299f11b44d8d3a588234adbe01112126010bd96d9139c3ba7b3badd9829261c3 \
+    --hash=sha256:2b24ec55fad43e476905eceaf14f41f6478780b870eda5d08b4d6de9a60b65b4 \
+    --hash=sha256:2d374d70fdc36f5863b84e54775452f68639bc862918602d028f89310a034ab0 \
+    --hash=sha256:2d9f0606baaec5dd54cb99667fcf85183a7477f3766fbddbe3f385e7fc253299 \
+    --hash=sha256:2e7ba4c9377e48fb7b20dedbd473cbcbc13e72e1826917c185157a137dac9df2 \
+    --hash=sha256:2f0a6423295a0d282d00e8701fe763eeefba8037e984ad5de44aa349002562ac \
+    --hash=sha256:327828786da2006085a4d1feb2594de6f6d26f8af48b81eb1ae950c788d97f61 \
+    --hash=sha256:380e6c38ef692b8fd5a0f6d1fa8774d81ebc08cfbd624b1bca62a4d4af2f9931 \
+    --hash=sha256:3b74ff4767d3ef47ffe0cd1d89379dc4d828d4873e5528976ced3b44fe5b0a21 \
+    --hash=sha256:3e844be8d536afa129366d9af76ed7cb8dfefec99f5f1c9e4f8ae542279a6dc3 \
+    --hash=sha256:459e81c2fb920b5f5df744262d1498ec2c8081acdcfe18181da44c50f51312f7 \
+    --hash=sha256:46ddf6e0b975cd680eb83318aa1d321cb2bf8d288d50f1754526230fcf59ba96 \
+    --hash=sha256:482c122b72e3c5ec98f11457aeb436ae4aecca75de19b3d1de7cf88bc40db82f \
+    --hash=sha256:561c87fea99545ef7d692403c110b2f99dced6dff93056d6e04384ad3bc46243 \
+    --hash=sha256:578d00c9b7fccfa1745a44f4eddfdc99d723d157dad26764538fbdda37209857 \
+    --hash=sha256:58c8e9620eb82a189c6c40cb6b59b4e35b2ee68b1f2afa6597732a2b467d7e8f \
+    --hash=sha256:5b29beab10211a746f9846baa39275e80034e065460d99eb51e45c9a9495bcca \
+    --hash=sha256:5d1d42556b063d579cae59e37a38c61f4402b47d70c29f0ef15cee1acaa64488 \
+    --hash=sha256:5f236cb5999ccd23a0ab1bd219cfe0ee3e1c1b65aaf6dd3320e972f7ec3a39da \
+    --hash=sha256:62a91aefff3d11bf60e5956d340eb507a983a7ec802b19072bb989ce120cd948 \
+    --hash=sha256:64cc6e97f14cf8a275d79c5002281f3040c12e2e4220623b5759ea7f9868d6a5 \
+    --hash=sha256:6f4c9156c4d1eb490fe374fb294deeb7bc7eaccda50e23775b2354b6a6739934 \
+    --hash=sha256:7294e38f9aa2e9f05f765b28ffdc5d81378508ce6dadbe93f6d464a8c9594473 \
+    --hash=sha256:7615058aabad54416ddac99ade09a5510cf77039a3b903e94e8922f25ed203d7 \
+    --hash=sha256:7e48cdb8226644e2fbd0bdb0a0f87906a3db07087f4de77a1b1b1ccfd9e93685 \
+    --hash=sha256:7f63d176a81555984e91f2c84c2a574a61cab7111cc907e176f0f01538e9ff6e \
+    --hash=sha256:7f6595c852ca544aaeeb32d357e62c9c780eac69dcd34e40cae7b55bc4fb1147 \
+    --hash=sha256:7fac95714b09da9278a0b52e492466f773cfe37651cf467a83a1b659be24bf71 \
+    --hash=sha256:81713b70bea5c1386dc2f32a8f0dab4148a2928c7495c808c541ee0aae614d67 \
+    --hash=sha256:846dd2e1243407133d3195d2d7e4ceefcaa5f5bf7278f0a9bda00967e6326b04 \
+    --hash=sha256:84c063af19ef5130084db70ada40ce63a84f6c1ef4d3dbc34e5e8c4febb20822 \
+    --hash=sha256:881764d610e3269964fc4bb3c19bb6fce55422828e152b885609ec176b41cf11 \
+    --hash=sha256:8994b29c462de9a8fce2d591028b986dbbe1b32f3ad600b2d3e1c482c93abad6 \
+    --hash=sha256:8c79e9d7e3d8a32d4824250a9c6401194fb4c2ad9a0cec8f6a96e09a582c2cc0 \
+    --hash=sha256:8ee427208c675f1b6e344a1f89376a9613fc30b52646a04ac0c1f6587c7e46ec \
+    --hash=sha256:949681f68e0e3c25377462be4b658500e85ca24323d9619fdc41f68d46a1ffda \
+    --hash=sha256:9e275792097c9f7e80741c36de3b61917aebecc08a67ae62899b074566ff8556 \
+    --hash=sha256:9fb815155aac6bfa8d86184079652c9715c812d506b22cfa369196ef4e99d1b4 \
+    --hash=sha256:a2a64e62c7a0edd07c1c917b0586655f3362d2c2d37d474db1a509efb96fea1c \
+    --hash=sha256:a7ac5b4984c468ce4f4a553df281450df0a34aefae02e58d77a0847be8d1e11f \
+    --hash=sha256:aa46dce75078fceaf7cecac5817422febb4355fbdda440db55206e3bd288cfb8 \
+    --hash=sha256:ae3476e934b9d714aa8000d2e4c01eb2590eee10b9d8cd03e7983ad65dfbfcba \
+    --hash=sha256:b0341e6d9a0c0e3cdc65857ef518bb05b410dbd70d749a0d33ac0f39e81a4258 \
+    --hash=sha256:b40d1bf6e6f74f7c0a567a9e5e778bbd4699d1d3d2c0fe46f4b717eef9e96b95 \
+    --hash=sha256:b5c4804e4039f487e942c13381e6c27b4b4e66066d94ef1fae3f6ba8b953f383 \
+    --hash=sha256:b5d6a6c9602fd4598fa07e0389e19fe199ae96449008d8304bf5d47cb745462e \
+    --hash=sha256:b5f1ac7359e17efe0b6e5fec21de34145caef22b260e978336f325d5c84e6938 \
+    --hash=sha256:c0167540094838ee9093ef6cc2c69d0074bbf84a432b4995835e8e5a0d984374 \
+    --hash=sha256:c180ac742a083e109c1a18151f4dd8675f32679985a1c750d2ff806796165b55 \
+    --hash=sha256:c73df5b6e8fabe2ddb74876fb82d9dd44cbace0ca12e8861ce9155ad3c886139 \
+    --hash=sha256:c7e177c619342e407415d4f35dec63d2d134d951e24b5166afcdfd1362828e17 \
+    --hash=sha256:cbad927ea8ed814622305d842c93412cb47bd39a496ed0f96bfd42b922b4a217 \
+    --hash=sha256:cc353841428d56b683a123a813e6a686e07026d6b1c5757970a877195f880c2d \
+    --hash=sha256:cc7c92c1baa629cb03ecb0c3d12564f172218fb1739f54bf5f3881844daadc6d \
+    --hash=sha256:cc7d768260f4ba4ea01741c1b5fe3d3a6c70eb91c87f4c8761bbcce5181beafe \
+    --hash=sha256:d0eea830b591dbc68e030c86a9569826145df485b2b4554874b07fea1275a199 \
+    --hash=sha256:d216e5d9b8749563c7f2c6f7a0831057ec844c68b4c11cb10fc62d4fd373c26d \
+    --hash=sha256:d401f07261dc5aa36c2e4efc308548f6ae943bfff20fcadb0a07517a26b196d8 \
+    --hash=sha256:d6324274b4e0e2fa1b3eccb25997b1c9ed134ff61d296448ab8269f5ac068c4c \
+    --hash=sha256:d8a8b74d843c2638f3864a17d97a4acda58e40d3e44b6303b8cc3d3c44ae2d29 \
+    --hash=sha256:d9b6b28a57feb51605d6ae5e61a9044a31742db557a3b851a74c13bc61de5172 \
+    --hash=sha256:de599af166970d6a61accde358ec9ded821234cbbc8c6413acfec06056b8e860 \
+    --hash=sha256:e594b22688d5747b06e957f1ef822060cb5cb35b493066e33ceac0cf882188b7 \
+    --hash=sha256:e5b078134f48552c4d9527db2f7da0b5359abd49393cdf9794017baec7506170 \
+    --hash=sha256:eb6dce402734575e1a8cc0bb1509afca508a400a57ce13d306ea2c663bad1138 \
+    --hash=sha256:f1790a4b1e8e8e028c391175433b9c8122c39b46e1663228158e61e6f915bf06 \
+    --hash=sha256:f5efe0661b9fcd6246f27957f6ae1c0eb29bc60552820f01e970b4996e016004 \
+    --hash=sha256:f9cbfbc5faca235fbdf531b93aa0f9f005ec7d267d9d738761a4d42b744ea159 \
+    --hash=sha256:fbea1751729afe607d84acfd01efd95e3b31db148a181a441984ce9b3d3469da \
+    --hash=sha256:fca4b4307ebe9c3ec77a084da3a9d1999d164693d16492ca2b64594340999988 \
+    --hash=sha256:ff5c6771c7e3511a06555afa317879b7db8d640137ba55d6ab0d0c50425cab75
+    # via aiohttp
+zarr==2.18.3 \
+    --hash=sha256:2580d8cb6dd84621771a10d31c4d777dca8a27706a1a89b29f42d2d37e2df5ce \
+    --hash=sha256:b1f7dfd2496f436745cdd4c7bcf8d3b4bc1dceef5fdd0d589c87130d842496dd
+    # via xarray-tensorstore
 zipp==3.20.2 \
     --hash=sha256:a817ac80d6cf4b23bf7f2828b7cabf326f15a001bea8b1f9b49631780ba28350 \
     --hash=sha256:bc9eb26f4506fda01b81bcde0ca78103b6e62f991b381fec825435c836edbc29
diff --git a/builddeps/requirements_lock_3_12.txt b/builddeps/requirements_lock_3_12.txt
index e3840978..7ac01a50 100644
--- a/builddeps/requirements_lock_3_12.txt
+++ b/builddeps/requirements_lock_3_12.txt
@@ -16,10 +16,121 @@ absl-py==2.1.0 \
     #   ml-collections
     #   optax
     #   orbax-checkpoint
+    #   tensorflow-probability
+aiohappyeyeballs==2.4.3 \
+    --hash=sha256:75cf88a15106a5002a8eb1dab212525c00d1f4c0fa96e551c9fbe6f09a621586 \
+    --hash=sha256:8a7a83727b2756f394ab2895ea0765a0a8c475e3c71e98d43d76f22b4b435572
+    # via aiohttp
+aiohttp==3.10.10 \
+    --hash=sha256:007ec22fbc573e5eb2fb7dec4198ef8f6bf2fe4ce20020798b2eb5d0abda6138 \
+    --hash=sha256:00819de9e45d42584bed046314c40ea7e9aea95411b38971082cad449392b08c \
+    --hash=sha256:01948b1d570f83ee7bbf5a60ea2375a89dfb09fd419170e7f5af029510033d24 \
+    --hash=sha256:038f514fe39e235e9fef6717fbf944057bfa24f9b3db9ee551a7ecf584b5b480 \
+    --hash=sha256:03a42ac7895406220124c88911ebee31ba8b2d24c98507f4a8bf826b2937c7f2 \
+    --hash=sha256:05646ebe6b94cc93407b3bf34b9eb26c20722384d068eb7339de802154d61bc5 \
+    --hash=sha256:0631dd7c9f0822cc61c88586ca76d5b5ada26538097d0f1df510b082bad3411a \
+    --hash=sha256:0b00807e2605f16e1e198f33a53ce3c4523114059b0c09c337209ae55e3823a8 \
+    --hash=sha256:0e1b370d8007c4ae31ee6db7f9a2fe801a42b146cec80a86766e7ad5c4a259cf \
+    --hash=sha256:15ecd889a709b0080f02721255b3f80bb261c2293d3c748151274dfea93ac871 \
+    --hash=sha256:1b66ccafef7336a1e1f0e389901f60c1d920102315a56df85e49552308fc0486 \
+    --hash=sha256:1bbb122c557a16fafc10354b9d99ebf2f2808a660d78202f10ba9d50786384b9 \
+    --hash=sha256:1eb89d3d29adaf533588f209768a9c02e44e4baf832b08118749c5fad191781d \
+    --hash=sha256:258c5dd01afc10015866114e210fb7365f0d02d9d059c3c3415382ab633fcbcb \
+    --hash=sha256:2609e9ab08474702cc67b7702dbb8a80e392c54613ebe80db7e8dbdb79837c68 \
+    --hash=sha256:274cfa632350225ce3fdeb318c23b4a10ec25c0e2c880eff951a3842cf358ac1 \
+    --hash=sha256:28529e08fde6f12eba8677f5a8608500ed33c086f974de68cc65ab218713a59d \
+    --hash=sha256:2b606353da03edcc71130b52388d25f9a30a126e04caef1fd637e31683033abd \
+    --hash=sha256:30ca7c3b94708a9d7ae76ff281b2f47d8eaf2579cd05971b5dc681db8caac6e1 \
+    --hash=sha256:333cf6cf8e65f6a1e06e9eb3e643a0c515bb850d470902274239fea02033e9a8 \
+    --hash=sha256:3455522392fb15ff549d92fbf4b73b559d5e43dc522588f7eb3e54c3f38beee7 \
+    --hash=sha256:362f641f9071e5f3ee6f8e7d37d5ed0d95aae656adf4ef578313ee585b585959 \
+    --hash=sha256:3bcd391d083f636c06a68715e69467963d1f9600f85ef556ea82e9ef25f043f7 \
+    --hash=sha256:3dffb610a30d643983aeb185ce134f97f290f8935f0abccdd32c77bed9388b42 \
+    --hash=sha256:3fe407bf93533a6fa82dece0e74dbcaaf5d684e5a51862887f9eaebe6372cd79 \
+    --hash=sha256:413251f6fcf552a33c981c4709a6bba37b12710982fec8e558ae944bfb2abd38 \
+    --hash=sha256:438cd072f75bb6612f2aca29f8bd7cdf6e35e8f160bc312e49fbecab77c99e3a \
+    --hash=sha256:4470c73c12cd9109db8277287d11f9dd98f77fc54155fc71a7738a83ffcc8ea8 \
+    --hash=sha256:45c3b868724137f713a38376fef8120c166d1eadd50da1855c112fe97954aed8 \
+    --hash=sha256:486f7aabfa292719a2753c016cc3a8f8172965cabb3ea2e7f7436c7f5a22a151 \
+    --hash=sha256:4f05e9727ce409358baa615dbeb9b969db94324a79b5a5cea45d39bdb01d82e6 \
+    --hash=sha256:50aed5155f819873d23520919e16703fc8925e509abbb1a1491b0087d1cd969e \
+    --hash=sha256:50edbcad60d8f0e3eccc68da67f37268b5144ecc34d59f27a02f9611c1d4eec7 \
+    --hash=sha256:54ca74df1be3c7ca1cf7f4c971c79c2daf48d9aa65dea1a662ae18926f5bc8ce \
+    --hash=sha256:578a4b875af3e0daaf1ac6fa983d93e0bbfec3ead753b6d6f33d467100cdc67b \
+    --hash=sha256:597a079284b7ee65ee102bc3a6ea226a37d2b96d0418cc9047490f231dc09fe8 \
+    --hash=sha256:59bb3c54aa420521dc4ce3cc2c3fe2ad82adf7b09403fa1f48ae45c0cbde6628 \
+    --hash=sha256:5c6a5b8c7926ba5d8545c7dd22961a107526562da31a7a32fa2456baf040939f \
+    --hash=sha256:64f6c17757251e2b8d885d728b6433d9d970573586a78b78ba8929b0f41d045a \
+    --hash=sha256:679abe5d3858b33c2cf74faec299fda60ea9de62916e8b67e625d65bf069a3b7 \
+    --hash=sha256:741a46d58677d8c733175d7e5aa618d277cd9d880301a380fd296975a9cdd7bc \
+    --hash=sha256:7789050d9e5d0c309c706953e5e8876e38662d57d45f936902e176d19f1c58ab \
+    --hash=sha256:77abf6665ae54000b98b3c742bc6ea1d1fb31c394bcabf8b5d2c1ac3ebfe7f3b \
+    --hash=sha256:79019094f87c9fb44f8d769e41dbb664d6e8fcfd62f665ccce36762deaa0e911 \
+    --hash=sha256:7b06b7843929e41a94ea09eb1ce3927865387e3e23ebe108e0d0d09b08d25be9 \
+    --hash=sha256:7e338c0523d024fad378b376a79faff37fafb3c001872a618cde1d322400a572 \
+    --hash=sha256:7ea7ffc6d6d6f8a11e6f40091a1040995cdff02cfc9ba4c2f30a516cb2633554 \
+    --hash=sha256:8105fd8a890df77b76dd3054cddf01a879fc13e8af576805d667e0fa0224c35d \
+    --hash=sha256:84afcdea18eda514c25bc68b9af2a2b1adea7c08899175a51fe7c4fb6d551257 \
+    --hash=sha256:9294bbb581f92770e6ed5c19559e1e99255e4ca604a22c5c6397b2f9dd3ee42c \
+    --hash=sha256:93429602396f3383a797a2a70e5f1de5df8e35535d7806c9f91df06f297e109b \
+    --hash=sha256:9627cc1a10c8c409b5822a92d57a77f383b554463d1884008e051c32ab1b3742 \
+    --hash=sha256:998f3bd3cfc95e9424a6acd7840cbdd39e45bc09ef87533c006f94ac47296090 \
+    --hash=sha256:9c72109213eb9d3874f7ac8c0c5fa90e072d678e117d9061c06e30c85b4cf0e6 \
+    --hash=sha256:9fc1500fd2a952c5c8e3b29aaf7e3cc6e27e9cfc0a8819b3bce48cc1b849e4cc \
+    --hash=sha256:a3f00003de6eba42d6e94fabb4125600d6e484846dbf90ea8e48a800430cc142 \
+    --hash=sha256:a45d85cf20b5e0d0aa5a8dca27cce8eddef3292bc29d72dcad1641f4ed50aa16 \
+    --hash=sha256:a7d8d14fe962153fc681f6366bdec33d4356f98a3e3567782aac1b6e0e40109a \
+    --hash=sha256:a8fa23fe62c436ccf23ff930149c047f060c7126eae3ccea005f0483f27b2e28 \
+    --hash=sha256:aa6658732517ddabe22c9036479eabce6036655ba87a0224c612e1ae6af2087e \
+    --hash=sha256:aafc8ee9b742ce75044ae9a4d3e60e3d918d15a4c2e08a6c3c3e38fa59b92d94 \
+    --hash=sha256:ab5a5a0c7a7991d90446a198689c0535be89bbd6b410a1f9a66688f0880ec026 \
+    --hash=sha256:acd48d5b80ee80f9432a165c0ac8cbf9253eaddb6113269a5e18699b33958dbb \
+    --hash=sha256:ad7593bb24b2ab09e65e8a1d385606f0f47c65b5a2ae6c551db67d6653e78c28 \
+    --hash=sha256:baa42524a82f75303f714108fea528ccacf0386af429b69fff141ffef1c534f9 \
+    --hash=sha256:bdfcf6443637c148c4e1a20c48c566aa694fa5e288d34b20fcdc58507882fed3 \
+    --hash=sha256:be7443669ae9c016b71f402e43208e13ddf00912f47f623ee5994e12fc7d4b3f \
+    --hash=sha256:c02a30b904282777d872266b87b20ed8cc0d1501855e27f831320f471d54d983 \
+    --hash=sha256:c1277cd707c465cd09572a774559a3cc7c7a28802eb3a2a9472588f062097205 \
+    --hash=sha256:c30a0eafc89d28e7f959281b58198a9fa5e99405f716c0289b7892ca345fe45f \
+    --hash=sha256:c5ce2ce7c997e1971b7184ee37deb6ea9922ef5163c6ee5aa3c274b05f9e12fa \
+    --hash=sha256:c823bc3971c44ab93e611ab1a46b1eafeae474c0c844aff4b7474287b75fe49c \
+    --hash=sha256:ce0cdc074d540265bfeb31336e678b4e37316849d13b308607efa527e981f5c2 \
+    --hash=sha256:d1720b4f14c78a3089562b8875b53e36b51c97c51adc53325a69b79b4b48ebcb \
+    --hash=sha256:d183cf9c797a5291e8301790ed6d053480ed94070637bfaad914dd38b0981f67 \
+    --hash=sha256:d9010c31cd6fa59438da4e58a7f19e4753f7f264300cd152e7f90d4602449762 \
+    --hash=sha256:d9e5e4a85bdb56d224f412d9c98ae4cbd032cc4f3161818f692cd81766eee65a \
+    --hash=sha256:da1dee8948d2137bb51fbb8a53cce6b1bcc86003c6b42565f008438b806cccd8 \
+    --hash=sha256:df9270660711670e68803107d55c2b5949c2e0f2e4896da176e1ecfc068b974a \
+    --hash=sha256:e00e3505cd80440f6c98c6d69269dcc2a119f86ad0a9fd70bccc59504bebd68a \
+    --hash=sha256:e48d5021a84d341bcaf95c8460b152cfbad770d28e5fe14a768988c461b821bc \
+    --hash=sha256:e7f8b04d83483577fd9200461b057c9f14ced334dcb053090cea1da9c8321a91 \
+    --hash=sha256:edfe3341033a6b53a5c522c802deb2079eee5cbfbb0af032a55064bd65c73a23 \
+    --hash=sha256:ef9c33cc5cbca35808f6c74be11eb7f5f6b14d2311be84a15b594bd3e58b5527 \
+    --hash=sha256:f2d4324a98062be0525d16f768a03e0bbb3b9fe301ceee99611dc9a7953124e6 \
+    --hash=sha256:f3935f82f6f4a3820270842e90456ebad3af15810cf65932bd24da4463bc0a4c \
+    --hash=sha256:f614ab0c76397661b90b6851a030004dac502e48260ea10f2441abd2207fbcc7 \
+    --hash=sha256:f7db54c7914cc99d901d93a34704833568d86c20925b2762f9fa779f9cd2e70f \
+    --hash=sha256:fbc6264158392bad9df19537e872d476f7c57adf718944cc1e4495cbabf38e2a \
+    --hash=sha256:fe2fb38c2ed905a2582948e2de560675e9dfbee94c6d5ccdb1301c6d0a5bf092 \
+    --hash=sha256:ffe595f10566f8276b76dc3a11ae4bb7eba1aac8ddd75811736a15b0d5311414
+    # via gcsfs
+aiosignal==1.3.1 \
+    --hash=sha256:54cd96e15e1649b75d6c87526a6ff0b6c1b0dd3459f43d9ca11d48c339b68cfc \
+    --hash=sha256:f8376fb07dd1e86a584e4fcdec80b36b7f81aac666ebc724e2c090300dd83b17
+    # via aiohttp
+asciitree==0.3.3 \
+    --hash=sha256:4aa4b9b649f85e3fcb343363d97564aa1fb62e249677f2e18a96765145cc0f6e
+    # via zarr
 attrs==24.2.0 \
     --hash=sha256:5cfb1b9148b5b086569baec03f20d7b6bf3bcacc9a42bebf87ffaaca362f6346 \
     --hash=sha256:81921eb96de3191c8258c199618104dd27ac608d9366f5e35d011eae1867ede2
-    # via e3nn-jax
+    # via
+    #   aiohttp
+    #   e3nn-jax
+cachetools==5.5.0 \
+    --hash=sha256:02134e8439cdc2ffb62023ce1debca2944c3f289d66bb17ead3ab3dede74b292 \
+    --hash=sha256:2cc24fb4cbe39633fb7badd9db9ca6295d766d9c2995f245725a46715d050f2a
+    # via google-auth
 certifi==2024.7.4 \
     --hash=sha256:5a1e7645bc0ec61a09e26c36f6106dd4cf40c6db3a1fb6352b0244e7fb057c7b \
     --hash=sha256:c198e21b1289c2ab85ee4e67bb4b4ef3ead0892059901a8d5b622f24a1101e90
@@ -120,14 +231,78 @@ chex==0.1.87 \
     --hash=sha256:0096d89cc8d898bb521ef4bfbf5c24549022b0e5b301f529ab57238896fe6c5d \
     --hash=sha256:ce536475661fd96d21be0c1728ecdbedd03f8ff950c662dfc338c92ea782cb16
     # via optax
+cloudpickle==3.1.0 \
+    --hash=sha256:81a929b6e3c7335c863c771d673d105f02efdb89dfaba0c90495d1c64796601b \
+    --hash=sha256:fe11acda67f61aaaec473e3afe030feb131d78a43461b718185363384f1ba12e
+    # via tensorflow-probability
 contextlib2==21.6.0 \
     --hash=sha256:3fbdb64466afd23abaf6c977627b75b6139a5a3e8ce38405c5b413aed7a0471f \
     --hash=sha256:ab1e2bfe1d01d968e1b7e8d9023bc51ef3509bba217bb730cee3827e1ee82869
     # via ml-collections
+decorator==5.1.1 \
+    --hash=sha256:637996211036b6385ef91435e4fae22989472f9d571faba8927ba8253acbc330 \
+    --hash=sha256:b8c3f85900b9dc423225913c5aace94729fe1fa9763b38939a95226f02d37186
+    # via
+    #   gcsfs
+    #   tensorflow-probability
+dinosaur-dycore==1.0.0 \
+    --hash=sha256:0ea3807a83d0c5220ff7a1cd548c11d82091d455d17d8f7697b985b6283d2ef9 \
+    --hash=sha256:3505d59eb9d90ad58572f93fa87db5536b9df0663f95701d50c8ccb027e767e1
+    # via neuralgcm
 dm-haiku==0.0.12 \
     --hash=sha256:7448a43a6486bff95253f84e18eacc607d9c1256592573117a9d1d23e2780706 \
     --hash=sha256:ba0b3acf71433156737fe342c486da11727e5e6c9e054245f4f9b8f0b53eb608
-    # via jax-md
+    # via
+    #   jax-md
+    #   neuralgcm
+dm-tree==0.1.8 \
+    --hash=sha256:054b461f8176f4bce7a21f7b1870f873a1ced3bdbe1282c816c550bb43c71fa6 \
+    --hash=sha256:09964470f76a5201aff2e8f9b26842976de7889300676f927930f6285e256760 \
+    --hash=sha256:0d3172394079a86c3a759179c65f64c48d1a42b89495fcf38976d11cc3bb952c \
+    --hash=sha256:0e9620ccf06393eb6b613b5e366469304622d4ea96ae6540b28a33840e6c89cf \
+    --hash=sha256:0fcaabbb14e7980377439e7140bd05552739ca5e515ecb3119f234acee4b9430 \
+    --hash=sha256:1607ce49aa42f010d1e5e616d92ce899d66835d4d8bea49679582435285515de \
+    --hash=sha256:181c35521d480d0365f39300542cb6cd7fd2b77351bb43d7acfda15aef63b317 \
+    --hash=sha256:1d7c26e431fc93cc7e0cba867eb000db6a05f6f2b25af11ac4e9dada88fc5bca \
+    --hash=sha256:1fe962015b2fe1282892b28ebe962faed53c7f98d942da9a4625cbf27baef913 \
+    --hash=sha256:250b692fb75f45f02e2f58fbef9ab338904ef334b90557565621fa251df267cf \
+    --hash=sha256:2869228d9c619074de501a3c10dc7f07c75422f8fab36ecdcb859b6f1b1ec3ef \
+    --hash=sha256:28c52cbf4f8b3dbd0beaedf44f69fa85eec5e9dede612e08035e06ada6ec9426 \
+    --hash=sha256:2f7915660f59c09068e428613c480150180df1060561fd0d1470684ae7007bd1 \
+    --hash=sha256:343a4a4ebaa127451ff971254a4be4084eb4bdc0b2513c32b46f6f728fd03f9e \
+    --hash=sha256:35cc164a79336bfcfafb47e5f297898359123bbd3330c1967f0c4994f9cf9f60 \
+    --hash=sha256:378cc8ad93c5fe3590f405a309980721f021c790ca1bdf9b15bb1d59daec57f5 \
+    --hash=sha256:39070ba268c0491af9fe7a58644d99e8b4f2cde6e5884ba3380bddc84ed43d5f \
+    --hash=sha256:435227cf3c5dc63f4de054cf3d00183790bd9ead4c3623138c74dde7f67f521b \
+    --hash=sha256:5483dca4d7eb1a0d65fe86d3b6a53ae717face83c1f17e0887b1a4a64ae5c410 \
+    --hash=sha256:694c3654cfd2a81552c08ec66bb5c4a3d48fa292b9a181880fb081c36c5b9134 \
+    --hash=sha256:75c5d528bb992981c20793b6b453e91560784215dffb8a5440ba999753c14ceb \
+    --hash=sha256:803bfc53b4659f447ac694dbd04235f94a73ef7c1fd1e0df7c84ac41e0bc963b \
+    --hash=sha256:81fce77f22a302d7a5968aebdf4efafef4def7ce96528719a354e6990dcd49c7 \
+    --hash=sha256:83b7764de0d855338abefc6e3ee9fe40d301668310aa3baea3f778ff051f4393 \
+    --hash=sha256:8c60a7eadab64c2278861f56bca320b2720f163dca9d7558103c3b77f2416571 \
+    --hash=sha256:8ed3564abed97c806db122c2d3e1a2b64c74a63debe9903aad795167cc301368 \
+    --hash=sha256:94d3f0826311f45ee19b75f5b48c99466e4218a0489e81c0f0167bda50cacf22 \
+    --hash=sha256:96a548a406a6fb15fe58f6a30a57ff2f2aafbf25f05afab00c8f5e5977b6c715 \
+    --hash=sha256:a5d819c38c03f0bb5b3b3703c60e4b170355a0fc6b5819325bf3d4ceb3ae7e80 \
+    --hash=sha256:ad16ceba90a56ec47cf45b21856d14962ac314787975ef786efb5e6e9ca75ec7 \
+    --hash=sha256:af4b3d372f2477dcd89a6e717e4a575ca35ccc20cc4454a8a4b6f8838a00672d \
+    --hash=sha256:b095ba4f8ca1ba19350fd53cf1f8f3eb0bd406aa28af64a6dfc86707b32a810a \
+    --hash=sha256:b9bd9b9ccb59409d33d51d84b7668010c04c2af7d4a371632874c1ca356cff3d \
+    --hash=sha256:b9f89a454e98806b44fe9d40ec9eee61f848388f7e79ac2371a55679bd5a3ac6 \
+    --hash=sha256:bb2d109f42190225112da899b9f3d46d0d5f26aef501c61e43529fe9322530b5 \
+    --hash=sha256:c0a94aba18a35457a1b5cd716fd7b46c5dafdc4cf7869b4bae665b91c4682a8e \
+    --hash=sha256:c5c8c12e3fda754ef6af94161bacdaeda816d941995fac415d6855c6c386af68 \
+    --hash=sha256:d1612fcaecd79023dbc6a6ae48d51a80beb5c385d6f3f6d71688e57bc8d07de8 \
+    --hash=sha256:d16e1f2a073604cfcc09f7131ae8d534674f43c3aef4c25742eae295bc60d04f \
+    --hash=sha256:d20f2faa3672b52e5013f4077117bfb99c4cfc0b445d3bde1584c34032b57436 \
+    --hash=sha256:d40fa4106ca6edc66760246a08f500ec0c85ef55c762fb4a363f6ee739ba02ee \
+    --hash=sha256:de287fabc464b8734be251e46e06aa9aa1001f34198da2b6ce07bd197172b9cb \
+    --hash=sha256:e4d714371bb08839e4e5e29024fc95832d9affe129825ef38836b143028bd144 \
+    --hash=sha256:ea9e59e0451e7d29aece402d9f908f2e2a80922bcde2ebfd5dcb07750fcbfee8 \
+    --hash=sha256:f7ac31b9aecccb2c6e1ab29706f6ded3eba0c2c69c770322c9c685929c3d6afb \
+    --hash=sha256:fa42a605d099ee7d41ba2b5fb75e21423951fd26e5d50583a00471238fb3021d
+    # via tensorflow-probability
 e3nn-jax==0.20.7 \
     --hash=sha256:7764edb99d19465e574a23857e5f97420b4d720f06ba8bc18a01f9f2b8301d1c \
     --hash=sha256:8a1b6f0c00feed565358fceec4c916a623412b6328945fce272f0c33e52b4bfc
@@ -142,12 +317,24 @@ etils[epath,epy]==1.9.4 \
     # via
     #   optax
     #   orbax-checkpoint
+fasteners==0.19 \
+    --hash=sha256:758819cb5d94cdedf4e836988b74de396ceacb8e2794d21f82d131fd9ee77237 \
+    --hash=sha256:b4f37c3ac52d8a445af3a66bce57b33b5e90b97c696b7b984f530cf8f0ded09c
+    # via zarr
 flax==0.9.0 \
     --hash=sha256:12cd8f7162165ddd56877fb1cd9a4fcb47a31569e4c5343eeb59a36369fa2cfe \
     --hash=sha256:8b7f361eed0f5324e81f9dc8d02ea53da5f993d7c2e37e7aa5b37d3f6331dd53
     # via
     #   dm-haiku
     #   jax-md
+flexcache==0.3 \
+    --hash=sha256:18743bd5a0621bfe2cf8d519e4c3bfdf57a269c15d1ced3fb4b64e0ff4600656 \
+    --hash=sha256:d43c9fea82336af6e0115e308d9d33a185390b8346a017564611f1466dcd2e32
+    # via pint
+flexparser==0.4 \
+    --hash=sha256:266d98905595be2ccc5da964fe0a2c3526fbbffdc45b65b3146d75db992ef6b2 \
+    --hash=sha256:3738b456192dcb3e15620f324c447721023c0293f6af9955b481e91d00179846
+    # via pint
 frozendict==2.4.4 \
     --hash=sha256:07c3a5dee8bbb84cba770e273cdbf2c87c8e035903af8f781292d72583416801 \
     --hash=sha256:12a342e439aef28ccec533f0253ea53d75fe9102bd6ea928ff530e76eac38906 \
@@ -183,10 +370,187 @@ frozendict==2.4.4 \
     --hash=sha256:e1b941132d79ce72d562a13341d38fc217bc1ee24d8c35a20d754e79ff99e038 \
     --hash=sha256:fefeb700bc7eb8b4c2dc48704e4221860d254c8989fb53488540bc44e44a1ac2
     # via jax-md
+frozenlist==1.5.0 \
+    --hash=sha256:000a77d6034fbad9b6bb880f7ec073027908f1b40254b5d6f26210d2dab1240e \
+    --hash=sha256:03d33c2ddbc1816237a67f66336616416e2bbb6beb306e5f890f2eb22b959cdf \
+    --hash=sha256:04a5c6babd5e8fb7d3c871dc8b321166b80e41b637c31a995ed844a6139942b6 \
+    --hash=sha256:0996c66760924da6e88922756d99b47512a71cfd45215f3570bf1e0b694c206a \
+    --hash=sha256:0cc974cc93d32c42e7b0f6cf242a6bd941c57c61b618e78b6c0a96cb72788c1d \
+    --hash=sha256:0f253985bb515ecd89629db13cb58d702035ecd8cfbca7d7a7e29a0e6d39af5f \
+    --hash=sha256:11aabdd62b8b9c4b84081a3c246506d1cddd2dd93ff0ad53ede5defec7886b28 \
+    --hash=sha256:12f78f98c2f1c2429d42e6a485f433722b0061d5c0b0139efa64f396efb5886b \
+    --hash=sha256:140228863501b44b809fb39ec56b5d4071f4d0aa6d216c19cbb08b8c5a7eadb9 \
+    --hash=sha256:1431d60b36d15cda188ea222033eec8e0eab488f39a272461f2e6d9e1a8e63c2 \
+    --hash=sha256:15538c0cbf0e4fa11d1e3a71f823524b0c46299aed6e10ebb4c2089abd8c3bec \
+    --hash=sha256:15b731db116ab3aedec558573c1a5eec78822b32292fe4f2f0345b7f697745c2 \
+    --hash=sha256:17dcc32fc7bda7ce5875435003220a457bcfa34ab7924a49a1c19f55b6ee185c \
+    --hash=sha256:1893f948bf6681733aaccf36c5232c231e3b5166d607c5fa77773611df6dc336 \
+    --hash=sha256:189f03b53e64144f90990d29a27ec4f7997d91ed3d01b51fa39d2dbe77540fd4 \
+    --hash=sha256:1a8ea951bbb6cacd492e3948b8da8c502a3f814f5d20935aae74b5df2b19cf3d \
+    --hash=sha256:1b96af8c582b94d381a1c1f51ffaedeb77c821c690ea5f01da3d70a487dd0a9b \
+    --hash=sha256:1e76bfbc72353269c44e0bc2cfe171900fbf7f722ad74c9a7b638052afe6a00c \
+    --hash=sha256:2150cc6305a2c2ab33299453e2968611dacb970d2283a14955923062c8d00b10 \
+    --hash=sha256:226d72559fa19babe2ccd920273e767c96a49b9d3d38badd7c91a0fdeda8ea08 \
+    --hash=sha256:237f6b23ee0f44066219dae14c70ae38a63f0440ce6750f868ee08775073f942 \
+    --hash=sha256:29d94c256679247b33a3dc96cce0f93cbc69c23bf75ff715919332fdbb6a32b8 \
+    --hash=sha256:2b5e23253bb709ef57a8e95e6ae48daa9ac5f265637529e4ce6b003a37b2621f \
+    --hash=sha256:2d0da8bbec082bf6bf18345b180958775363588678f64998c2b7609e34719b10 \
+    --hash=sha256:2f3f7a0fbc219fb4455264cae4d9f01ad41ae6ee8524500f381de64ffaa077d5 \
+    --hash=sha256:30c72000fbcc35b129cb09956836c7d7abf78ab5416595e4857d1cae8d6251a6 \
+    --hash=sha256:31115ba75889723431aa9a4e77d5f398f5cf976eea3bdf61749731f62d4a4a21 \
+    --hash=sha256:31a9ac2b38ab9b5a8933b693db4939764ad3f299fcaa931a3e605bc3460e693c \
+    --hash=sha256:366d8f93e3edfe5a918c874702f78faac300209a4d5bf38352b2c1bdc07a766d \
+    --hash=sha256:374ca2dabdccad8e2a76d40b1d037f5bd16824933bf7bcea3e59c891fd4a0923 \
+    --hash=sha256:44c49271a937625619e862baacbd037a7ef86dd1ee215afc298a417ff3270608 \
+    --hash=sha256:45e0896250900b5aa25180f9aec243e84e92ac84bd4a74d9ad4138ef3f5c97de \
+    --hash=sha256:498524025a5b8ba81695761d78c8dd7382ac0b052f34e66939c42df860b8ff17 \
+    --hash=sha256:50cf5e7ee9b98f22bdecbabf3800ae78ddcc26e4a435515fc72d97903e8488e0 \
+    --hash=sha256:52ef692a4bc60a6dd57f507429636c2af8b6046db8b31b18dac02cbc8f507f7f \
+    --hash=sha256:561eb1c9579d495fddb6da8959fd2a1fca2c6d060d4113f5844b433fc02f2641 \
+    --hash=sha256:5a3ba5f9a0dfed20337d3e966dc359784c9f96503674c2faf015f7fe8e96798c \
+    --hash=sha256:5b6a66c18b5b9dd261ca98dffcb826a525334b2f29e7caa54e182255c5f6a65a \
+    --hash=sha256:5c28f4b5dbef8a0d8aad0d4de24d1e9e981728628afaf4ea0792f5d0939372f0 \
+    --hash=sha256:5d7f5a50342475962eb18b740f3beecc685a15b52c91f7d975257e13e029eca9 \
+    --hash=sha256:6321899477db90bdeb9299ac3627a6a53c7399c8cd58d25da094007402b039ab \
+    --hash=sha256:6482a5851f5d72767fbd0e507e80737f9c8646ae7fd303def99bfe813f76cf7f \
+    --hash=sha256:666534d15ba8f0fda3f53969117383d5dc021266b3c1a42c9ec4855e4b58b9d3 \
+    --hash=sha256:683173d371daad49cffb8309779e886e59c2f369430ad28fe715f66d08d4ab1a \
+    --hash=sha256:6e9080bb2fb195a046e5177f10d9d82b8a204c0736a97a153c2466127de87784 \
+    --hash=sha256:73f2e31ea8dd7df61a359b731716018c2be196e5bb3b74ddba107f694fbd7604 \
+    --hash=sha256:7437601c4d89d070eac8323f121fcf25f88674627505334654fd027b091db09d \
+    --hash=sha256:76e4753701248476e6286f2ef492af900ea67d9706a0155335a40ea21bf3b2f5 \
+    --hash=sha256:7707a25d6a77f5d27ea7dc7d1fc608aa0a478193823f88511ef5e6b8a48f9d03 \
+    --hash=sha256:7948140d9f8ece1745be806f2bfdf390127cf1a763b925c4a805c603df5e697e \
+    --hash=sha256:7a1a048f9215c90973402e26c01d1cff8a209e1f1b53f72b95c13db61b00f953 \
+    --hash=sha256:7d57d8f702221405a9d9b40f9da8ac2e4a1a8b5285aac6100f3393675f0a85ee \
+    --hash=sha256:7f3c8c1dacd037df16e85227bac13cca58c30da836c6f936ba1df0c05d046d8d \
+    --hash=sha256:81d5af29e61b9c8348e876d442253723928dce6433e0e76cd925cd83f1b4b817 \
+    --hash=sha256:828afae9f17e6de596825cf4228ff28fbdf6065974e5ac1410cecc22f699d2b3 \
+    --hash=sha256:87f724d055eb4785d9be84e9ebf0f24e392ddfad00b3fe036e43f489fafc9039 \
+    --hash=sha256:8969190d709e7c48ea386db202d708eb94bdb29207a1f269bab1196ce0dcca1f \
+    --hash=sha256:90646abbc7a5d5c7c19461d2e3eeb76eb0b204919e6ece342feb6032c9325ae9 \
+    --hash=sha256:91d6c171862df0a6c61479d9724f22efb6109111017c87567cfeb7b5d1449fdf \
+    --hash=sha256:9272fa73ca71266702c4c3e2d4a28553ea03418e591e377a03b8e3659d94fa76 \
+    --hash=sha256:92b5278ed9d50fe610185ecd23c55d8b307d75ca18e94c0e7de328089ac5dcba \
+    --hash=sha256:97160e245ea33d8609cd2b8fd997c850b56db147a304a262abc2b3be021a9171 \
+    --hash=sha256:977701c081c0241d0955c9586ffdd9ce44f7a7795df39b9151cd9a6fd0ce4cfb \
+    --hash=sha256:9b7dc0c4338e6b8b091e8faf0db3168a37101943e687f373dce00959583f7439 \
+    --hash=sha256:9b93d7aaa36c966fa42efcaf716e6b3900438632a626fb09c049f6a2f09fc631 \
+    --hash=sha256:9bbcdfaf4af7ce002694a4e10a0159d5a8d20056a12b05b45cea944a4953f972 \
+    --hash=sha256:9c2623347b933fcb9095841f1cc5d4ff0b278addd743e0e966cb3d460278840d \
+    --hash=sha256:a2fe128eb4edeabe11896cb6af88fca5346059f6c8d807e3b910069f39157869 \
+    --hash=sha256:a72b7a6e3cd2725eff67cd64c8f13335ee18fc3c7befc05aed043d24c7b9ccb9 \
+    --hash=sha256:a9fe0f1c29ba24ba6ff6abf688cb0b7cf1efab6b6aa6adc55441773c252f7411 \
+    --hash=sha256:b97f7b575ab4a8af9b7bc1d2ef7f29d3afee2226bd03ca3875c16451ad5a7723 \
+    --hash=sha256:bdac3c7d9b705d253b2ce370fde941836a5f8b3c5c2b8fd70940a3ea3af7f4f2 \
+    --hash=sha256:c03eff4a41bd4e38415cbed054bbaff4a075b093e2394b6915dca34a40d1e38b \
+    --hash=sha256:c16d2fa63e0800723139137d667e1056bee1a1cf7965153d2d104b62855e9b99 \
+    --hash=sha256:c1fac3e2ace2eb1052e9f7c7db480818371134410e1f5c55d65e8f3ac6d1407e \
+    --hash=sha256:ce3aa154c452d2467487765e3adc730a8c153af77ad84096bc19ce19a2400840 \
+    --hash=sha256:cee6798eaf8b1416ef6909b06f7dc04b60755206bddc599f52232606e18179d3 \
+    --hash=sha256:d1b3eb7b05ea246510b43a7e53ed1653e55c2121019a97e60cad7efb881a97bb \
+    --hash=sha256:d994863bba198a4a518b467bb971c56e1db3f180a25c6cf7bb1949c267f748c3 \
+    --hash=sha256:dd47a5181ce5fcb463b5d9e17ecfdb02b678cca31280639255ce9d0e5aa67af0 \
+    --hash=sha256:dd94994fc91a6177bfaafd7d9fd951bc8689b0a98168aa26b5f543868548d3ca \
+    --hash=sha256:de537c11e4aa01d37db0d403b57bd6f0546e71a82347a97c6a9f0dcc532b3a45 \
+    --hash=sha256:df6e2f325bfee1f49f81aaac97d2aa757c7646534a06f8f577ce184afe2f0a9e \
+    --hash=sha256:e66cc454f97053b79c2ab09c17fbe3c825ea6b4de20baf1be28919460dd7877f \
+    --hash=sha256:e79225373c317ff1e35f210dd5f1344ff31066ba8067c307ab60254cd3a78ad5 \
+    --hash=sha256:f1577515d35ed5649d52ab4319db757bb881ce3b2b796d7283e6634d99ace307 \
+    --hash=sha256:f1e6540b7fa044eee0bb5111ada694cf3dc15f2b0347ca125ee9ca984d5e9e6e \
+    --hash=sha256:f2ac49a9bedb996086057b75bf93538240538c6d9b38e57c82d51f75a73409d2 \
+    --hash=sha256:f47c9c9028f55a04ac254346e92977bf0f166c483c74b4232bee19a6697e4778 \
+    --hash=sha256:f5f9da7f5dbc00a604fe74aa02ae7c98bcede8a3b8b9666f9f86fc13993bc71a \
+    --hash=sha256:fd74520371c3c4175142d02a976aee0b4cb4a7cc912a60586ffd8d5929979b30 \
+    --hash=sha256:feeb64bc9bcc6b45c6311c9e9b99406660a9c05ca8a5b30d14a78555088b0b3a
+    # via
+    #   aiohttp
+    #   aiosignal
 fsspec==2024.9.0 \
     --hash=sha256:4b0afb90c2f21832df142f292649035d80b421f60a9e1c027802e5a0da2b04e8 \
     --hash=sha256:a0947d552d8a6efa72cc2c730b12c41d043509156966cca4fb157b0f2a0c574b
-    # via etils
+    # via
+    #   dinosaur-dycore
+    #   etils
+    #   gcsfs
+gast==0.6.0 \
+    --hash=sha256:52b182313f7330389f72b069ba00f174cfe2a06411099547288839c6cbafbd54 \
+    --hash=sha256:88fc5300d32c7ac6ca7b515310862f71e6fdf2c029bbec7c66c0f5dd47b6b1fb
+    # via tensorflow-probability
+gcsfs==2024.9.0.post1 \
+    --hash=sha256:7ca70ee9d7c7dbce1a3e36b4883e14102c2d7b4284f49e242843a437bc684684 \
+    --hash=sha256:f3ab9d3bedc45da8cf40baed7c3a1e1694e8f599160d9138d78f0ef25e4a3ca1
+    # via -r builddeps/test-requirements.txt
+gin-config==0.5.0 \
+    --hash=sha256:0c6ea5026ded927c8c93c990b01c695257c1df446e45e549a158cfbc79e19ed6 \
+    --hash=sha256:bddb7ca221ea2b46cdb59321e79fecf02d6e3b728906047fcd4076c297609fd6
+    # via neuralgcm
+google-api-core==2.22.0 \
+    --hash=sha256:26f8d76b96477db42b55fd02a33aae4a42ec8b86b98b94969b7333a2c828bf35 \
+    --hash=sha256:a6652b6bd51303902494998626653671703c420f6f4c88cfd3f50ed723e9d021
+    # via
+    #   google-cloud-core
+    #   google-cloud-storage
+google-auth==2.36.0 \
+    --hash=sha256:51a15d47028b66fd36e5c64a82d2d57480075bccc7da37cde257fc94177a61fb \
+    --hash=sha256:545e9618f2df0bcbb7dcbc45a546485b1212624716975a1ea5ae8149ce769ab1
+    # via
+    #   gcsfs
+    #   google-api-core
+    #   google-auth-oauthlib
+    #   google-cloud-core
+    #   google-cloud-storage
+google-auth-oauthlib==1.2.1 \
+    --hash=sha256:2d58a27262d55aa1b87678c3ba7142a080098cbc2024f903c62355deb235d91f \
+    --hash=sha256:afd0cad092a2eaa53cd8e8298557d6de1034c6cb4a740500b5357b648af97263
+    # via gcsfs
+google-cloud-core==2.4.1 \
+    --hash=sha256:9b7749272a812bde58fff28868d0c5e2f585b82f37e09a1f6ed2d4d10f134073 \
+    --hash=sha256:a9e6a4422b9ac5c29f79a0ede9485473338e2ce78d91f2370c01e730eab22e61
+    # via google-cloud-storage
+google-cloud-storage==2.18.2 \
+    --hash=sha256:97a4d45c368b7d401ed48c4fdfe86e1e1cb96401c9e199e419d289e2c0370166 \
+    --hash=sha256:aaf7acd70cdad9f274d29332673fcab98708d0e1f4dceb5a5356aaef06af4d99
+    # via gcsfs
+google-crc32c==1.6.0 \
+    --hash=sha256:05e2d8c9a2f853ff116db9706b4a27350587f341eda835f46db3c0a8c8ce2f24 \
+    --hash=sha256:18e311c64008f1f1379158158bb3f0c8d72635b9eb4f9545f8cf990c5668e59d \
+    --hash=sha256:236c87a46cdf06384f614e9092b82c05f81bd34b80248021f729396a78e55d7e \
+    --hash=sha256:35834855408429cecf495cac67ccbab802de269e948e27478b1e47dfb6465e57 \
+    --hash=sha256:386122eeaaa76951a8196310432c5b0ef3b53590ef4c317ec7588ec554fec5d2 \
+    --hash=sha256:40b05ab32a5067525670880eb5d169529089a26fe35dce8891127aeddc1950e8 \
+    --hash=sha256:48abd62ca76a2cbe034542ed1b6aee851b6f28aaca4e6551b5599b6f3ef175cc \
+    --hash=sha256:50cf2a96da226dcbff8671233ecf37bf6e95de98b2a2ebadbfdf455e6d05df42 \
+    --hash=sha256:51c4f54dd8c6dfeb58d1df5e4f7f97df8abf17a36626a217f169893d1d7f3e9f \
+    --hash=sha256:5bcc90b34df28a4b38653c36bb5ada35671ad105c99cfe915fb5bed7ad6924aa \
+    --hash=sha256:62f6d4a29fea082ac4a3c9be5e415218255cf11684ac6ef5488eea0c9132689b \
+    --hash=sha256:6eceb6ad197656a1ff49ebfbbfa870678c75be4344feb35ac1edf694309413dc \
+    --hash=sha256:7aec8e88a3583515f9e0957fe4f5f6d8d4997e36d0f61624e70469771584c760 \
+    --hash=sha256:91ca8145b060679ec9176e6de4f89b07363d6805bd4760631ef254905503598d \
+    --hash=sha256:a184243544811e4a50d345838a883733461e67578959ac59964e43cca2c791e7 \
+    --hash=sha256:a9e4b426c3702f3cd23b933436487eb34e01e00327fac20c9aebb68ccf34117d \
+    --hash=sha256:bb0966e1c50d0ef5bc743312cc730b533491d60585a9a08f897274e57c3f70e0 \
+    --hash=sha256:bb8b3c75bd157010459b15222c3fd30577042a7060e29d42dabce449c087f2b3 \
+    --hash=sha256:bd5e7d2445d1a958c266bfa5d04c39932dc54093fa391736dbfdb0f1929c1fb3 \
+    --hash=sha256:c87d98c7c4a69066fd31701c4e10d178a648c2cac3452e62c6b24dc51f9fcc00 \
+    --hash=sha256:d2952396dc604544ea7476b33fe87faedc24d666fb0c2d5ac971a2b9576ab871 \
+    --hash=sha256:d8797406499f28b5ef791f339594b0b5fdedf54e203b5066675c406ba69d705c \
+    --hash=sha256:d9e9913f7bd69e093b81da4535ce27af842e7bf371cde42d1ae9e9bd382dc0e9 \
+    --hash=sha256:e2806553238cd076f0a55bddab37a532b53580e699ed8e5606d0de1f856b5205 \
+    --hash=sha256:ebab974b1687509e5c973b5c4b8b146683e101e102e17a86bd196ecaa4d099fc \
+    --hash=sha256:ed767bf4ba90104c1216b68111613f0d5926fb3780660ea1198fc469af410e9d \
+    --hash=sha256:f7a1fc29803712f80879b0806cb83ab24ce62fc8daf0569f2204a0cfd7f68ed4
+    # via
+    #   google-cloud-storage
+    #   google-resumable-media
+google-resumable-media==2.7.2 \
+    --hash=sha256:3ce7551e9fe6d99e9a126101d2536612bb73486721951e9562fee0f90c6ababa \
+    --hash=sha256:5280aed4629f2b60b847b0d42f9857fd4935c11af266744df33d8074cae92fe0
+    # via google-cloud-storage
+googleapis-common-protos==1.65.0 \
+    --hash=sha256:2972e6c496f435b92590fd54045060867f3fe9be2c82ab148fc8885035479a63 \
+    --hash=sha256:334a29d07cddc3aa01dee4988f9afd9b2916ee2ff49d6b757155dc0d197852c0
+    # via google-api-core
 humanize==4.11.0 \
     --hash=sha256:b53caaec8532bcb2fff70c8826f904c35943f8cecaca29d272d9df38092736c0 \
     --hash=sha256:e66f36020a2d5a974c504bd2555cf770621dbdbb6d82f94a6857c0b1ea2608be
@@ -194,7 +558,9 @@ humanize==4.11.0 \
 idna==3.8 \
     --hash=sha256:050b4e5baadcd44d760cedbd2b8e639f2ff89bbc7a5730fcc662954303377aac \
     --hash=sha256:d838c2c0ed6fced7693d5e8ab8e734d5f8fda53a039c0164afb0b82e771e3603
-    # via requests
+    # via
+    #   requests
+    #   yarl
 importlib-resources==6.4.5 \
     --hash=sha256:980862a1d16c9e147a59603677fa2aa5fd82b87f223b6cb870695bcfce830065 \
     --hash=sha256:ac29d5f956f01d5e4bb63102a5a19957f1b9175e45649977264a1416783bb717
@@ -206,12 +572,15 @@ jax==0.4.31 \
     #   -r builddeps/requirements.in
     #   -r builddeps/test-requirements.txt
     #   chex
+    #   dinosaur-dycore
     #   e3nn-jax
     #   flax
     #   jax-md
     #   jraph
+    #   neuralgcm
     #   optax
     #   orbax-checkpoint
+    #   tree-math
 jax-cuda12-pjrt==0.4.31 \
     --hash=sha256:3e77d1cfebeca06517254eb568f082037e1a2aa3ed8f63c543492ad8ab5a1585 \
     --hash=sha256:8961abb381d893a3c2392ad76ab2067a81f8f2514f3b47d2da3ac24283293fe0
@@ -247,19 +616,24 @@ jaxlib==0.4.31 \
     #   -r builddeps/requirements.in
     #   -r builddeps/test-requirements.txt
     #   chex
+    #   dinosaur-dycore
     #   e3nn-jax
     #   jax
     #   jax-md
     #   jraph
+    #   neuralgcm
     #   optax
     #   orbax-checkpoint
 jmp==0.0.4 \
     --hash=sha256:5dfeb0fd7c7a9f72a70fff0aab9d0cbfae32a809c02f4037ff3485ceb33e1730 \
     --hash=sha256:6aa7adbddf2bd574b28c7faf6e81a735eb11f53386447896909c6968dc36807d
     # via dm-haiku
+joblib==1.4.2 \
+    --hash=sha256:06d478d5674cbc267e7496a410ee875abd68e4340feff4490bcb7afb88060ae6 \
+    --hash=sha256:2382c5816b2636fbd20a09e0f4e9dad4736765fdfb7dca582943b9c1366b3f0e
+    # via scikit-learn
 jraph @ https://github.com/wsmoses/jraph/archive/b00d9a03db76c69a258a86df81638b9a2f28829f.tar.gz \
-    --hash=sha256:350fe37bf717f934f1f84fd3370a480b3178bfcb61dfa217c738971308c57625 \
-    --hash=sha256:c3ac3a0b224b344eb6d367e8bc312d95ea41bf825d01ea31b80dd8c22c0dd8b8
+    --hash=sha256:d996f47fff8c6ac8f5572e4220e3193f2542e9d3d6bc1defc16ac22dadd74dd4
     # via -r builddeps/test-requirements.txt
 markdown-it-py==3.0.0 \
     --hash=sha256:355216845c60bd96232cd8d8c40e8f9765cc86f46880e43a8fd22dc1a1a8cab1 \
@@ -366,10 +740,129 @@ msgpack==1.1.0 \
     # via
     #   flax
     #   orbax-checkpoint
+multidict==6.1.0 \
+    --hash=sha256:052e10d2d37810b99cc170b785945421141bf7bb7d2f8799d431e7db229c385f \
+    --hash=sha256:06809f4f0f7ab7ea2cabf9caca7d79c22c0758b58a71f9d32943ae13c7ace056 \
+    --hash=sha256:071120490b47aa997cca00666923a83f02c7fbb44f71cf7f136df753f7fa8761 \
+    --hash=sha256:0c3f390dc53279cbc8ba976e5f8035eab997829066756d811616b652b00a23a3 \
+    --hash=sha256:0e2b90b43e696f25c62656389d32236e049568b39320e2735d51f08fd362761b \
+    --hash=sha256:0e5f362e895bc5b9e67fe6e4ded2492d8124bdf817827f33c5b46c2fe3ffaca6 \
+    --hash=sha256:10524ebd769727ac77ef2278390fb0068d83f3acb7773792a5080f2b0abf7748 \
+    --hash=sha256:10a9b09aba0c5b48c53761b7c720aaaf7cf236d5fe394cd399c7ba662d5f9966 \
+    --hash=sha256:16e5f4bf4e603eb1fdd5d8180f1a25f30056f22e55ce51fb3d6ad4ab29f7d96f \
+    --hash=sha256:188215fc0aafb8e03341995e7c4797860181562380f81ed0a87ff455b70bf1f1 \
+    --hash=sha256:189f652a87e876098bbc67b4da1049afb5f5dfbaa310dd67c594b01c10388db6 \
+    --hash=sha256:1ca0083e80e791cffc6efce7660ad24af66c8d4079d2a750b29001b53ff59ada \
+    --hash=sha256:1e16bf3e5fc9f44632affb159d30a437bfe286ce9e02754759be5536b169b305 \
+    --hash=sha256:2090f6a85cafc5b2db085124d752757c9d251548cedabe9bd31afe6363e0aff2 \
+    --hash=sha256:20b9b5fbe0b88d0bdef2012ef7dee867f874b72528cf1d08f1d59b0e3850129d \
+    --hash=sha256:22ae2ebf9b0c69d206c003e2f6a914ea33f0a932d4aa16f236afc049d9958f4a \
+    --hash=sha256:22f3105d4fb15c8f57ff3959a58fcab6ce36814486500cd7485651230ad4d4ef \
+    --hash=sha256:23bfd518810af7de1116313ebd9092cb9aa629beb12f6ed631ad53356ed6b86c \
+    --hash=sha256:27e5fc84ccef8dfaabb09d82b7d179c7cf1a3fbc8a966f8274fcb4ab2eb4cadb \
+    --hash=sha256:3380252550e372e8511d49481bd836264c009adb826b23fefcc5dd3c69692f60 \
+    --hash=sha256:3702ea6872c5a2a4eeefa6ffd36b042e9773f05b1f37ae3ef7264b1163c2dcf6 \
+    --hash=sha256:37bb93b2178e02b7b618893990941900fd25b6b9ac0fa49931a40aecdf083fe4 \
+    --hash=sha256:3914f5aaa0f36d5d60e8ece6a308ee1c9784cd75ec8151062614657a114c4478 \
+    --hash=sha256:3a37ffb35399029b45c6cc33640a92bef403c9fd388acce75cdc88f58bd19a81 \
+    --hash=sha256:3c8b88a2ccf5493b6c8da9076fb151ba106960a2df90c2633f342f120751a9e7 \
+    --hash=sha256:3e97b5e938051226dc025ec80980c285b053ffb1e25a3db2a3aa3bc046bf7f56 \
+    --hash=sha256:3ec660d19bbc671e3a6443325f07263be452c453ac9e512f5eb935e7d4ac28b3 \
+    --hash=sha256:3efe2c2cb5763f2f1b275ad2bf7a287d3f7ebbef35648a9726e3b69284a4f3d6 \
+    --hash=sha256:483a6aea59cb89904e1ceabd2b47368b5600fb7de78a6e4a2c2987b2d256cf30 \
+    --hash=sha256:4867cafcbc6585e4b678876c489b9273b13e9fff9f6d6d66add5e15d11d926cb \
+    --hash=sha256:48e171e52d1c4d33888e529b999e5900356b9ae588c2f09a52dcefb158b27506 \
+    --hash=sha256:4a9cb68166a34117d6646c0023c7b759bf197bee5ad4272f420a0141d7eb03a0 \
+    --hash=sha256:4b820514bfc0b98a30e3d85462084779900347e4d49267f747ff54060cc33925 \
+    --hash=sha256:4e18b656c5e844539d506a0a06432274d7bd52a7487e6828c63a63d69185626c \
+    --hash=sha256:4e9f48f58c2c523d5a06faea47866cd35b32655c46b443f163d08c6d0ddb17d6 \
+    --hash=sha256:50b3a2710631848991d0bf7de077502e8994c804bb805aeb2925a981de58ec2e \
+    --hash=sha256:55b6d90641869892caa9ca42ff913f7ff1c5ece06474fbd32fb2cf6834726c95 \
+    --hash=sha256:57feec87371dbb3520da6192213c7d6fc892d5589a93db548331954de8248fd2 \
+    --hash=sha256:58130ecf8f7b8112cdb841486404f1282b9c86ccb30d3519faf301b2e5659133 \
+    --hash=sha256:5845c1fd4866bb5dd3125d89b90e57ed3138241540897de748cdf19de8a2fca2 \
+    --hash=sha256:59bfeae4b25ec05b34f1956eaa1cb38032282cd4dfabc5056d0a1ec4d696d3aa \
+    --hash=sha256:5b48204e8d955c47c55b72779802b219a39acc3ee3d0116d5080c388970b76e3 \
+    --hash=sha256:5c09fcfdccdd0b57867577b719c69e347a436b86cd83747f179dbf0cc0d4c1f3 \
+    --hash=sha256:6180c0ae073bddeb5a97a38c03f30c233e0a4d39cd86166251617d1bbd0af436 \
+    --hash=sha256:682b987361e5fd7a139ed565e30d81fd81e9629acc7d925a205366877d8c8657 \
+    --hash=sha256:6b5d83030255983181005e6cfbac1617ce9746b219bc2aad52201ad121226581 \
+    --hash=sha256:6bb5992037f7a9eff7991ebe4273ea7f51f1c1c511e6a2ce511d0e7bdb754492 \
+    --hash=sha256:73eae06aa53af2ea5270cc066dcaf02cc60d2994bbb2c4ef5764949257d10f43 \
+    --hash=sha256:76f364861c3bfc98cbbcbd402d83454ed9e01a5224bb3a28bf70002a230f73e2 \
+    --hash=sha256:820c661588bd01a0aa62a1283f20d2be4281b086f80dad9e955e690c75fb54a2 \
+    --hash=sha256:82176036e65644a6cc5bd619f65f6f19781e8ec2e5330f51aa9ada7504cc1926 \
+    --hash=sha256:87701f25a2352e5bf7454caa64757642734da9f6b11384c1f9d1a8e699758057 \
+    --hash=sha256:9079dfc6a70abe341f521f78405b8949f96db48da98aeb43f9907f342f627cdc \
+    --hash=sha256:90f8717cb649eea3504091e640a1b8568faad18bd4b9fcd692853a04475a4b80 \
+    --hash=sha256:957cf8e4b6e123a9eea554fa7ebc85674674b713551de587eb318a2df3e00255 \
+    --hash=sha256:99f826cbf970077383d7de805c0681799491cb939c25450b9b5b3ced03ca99f1 \
+    --hash=sha256:9f636b730f7e8cb19feb87094949ba54ee5357440b9658b2a32a5ce4bce53972 \
+    --hash=sha256:a114d03b938376557927ab23f1e950827c3b893ccb94b62fd95d430fd0e5cf53 \
+    --hash=sha256:a185f876e69897a6f3325c3f19f26a297fa058c5e456bfcff8015e9a27e83ae1 \
+    --hash=sha256:a7a9541cd308eed5e30318430a9c74d2132e9a8cb46b901326272d780bf2d423 \
+    --hash=sha256:aa466da5b15ccea564bdab9c89175c762bc12825f4659c11227f515cee76fa4a \
+    --hash=sha256:aaed8b0562be4a0876ee3b6946f6869b7bcdb571a5d1496683505944e268b160 \
+    --hash=sha256:ab7c4ceb38d91570a650dba194e1ca87c2b543488fe9309b4212694174fd539c \
+    --hash=sha256:ac10f4c2b9e770c4e393876e35a7046879d195cd123b4f116d299d442b335bcd \
+    --hash=sha256:b04772ed465fa3cc947db808fa306d79b43e896beb677a56fb2347ca1a49c1fa \
+    --hash=sha256:b1c416351ee6271b2f49b56ad7f308072f6f44b37118d69c2cad94f3fa8a40d5 \
+    --hash=sha256:b225d95519a5bf73860323e633a664b0d85ad3d5bede6d30d95b35d4dfe8805b \
+    --hash=sha256:b2f59caeaf7632cc633b5cf6fc449372b83bbdf0da4ae04d5be36118e46cc0aa \
+    --hash=sha256:b58c621844d55e71c1b7f7c498ce5aa6985d743a1a59034c57a905b3f153c1ef \
+    --hash=sha256:bf6bea52ec97e95560af5ae576bdac3aa3aae0b6758c6efa115236d9e07dae44 \
+    --hash=sha256:c08be4f460903e5a9d0f76818db3250f12e9c344e79314d1d570fc69d7f4eae4 \
+    --hash=sha256:c7053d3b0353a8b9de430a4f4b4268ac9a4fb3481af37dfe49825bf45ca24156 \
+    --hash=sha256:c943a53e9186688b45b323602298ab727d8865d8c9ee0b17f8d62d14b56f0753 \
+    --hash=sha256:ce2186a7df133a9c895dea3331ddc5ddad42cdd0d1ea2f0a51e5d161e4762f28 \
+    --hash=sha256:d093be959277cb7dee84b801eb1af388b6ad3ca6a6b6bf1ed7585895789d027d \
+    --hash=sha256:d094ddec350a2fb899fec68d8353c78233debde9b7d8b4beeafa70825f1c281a \
+    --hash=sha256:d1a9dd711d0877a1ece3d2e4fea11a8e75741ca21954c919406b44e7cf971304 \
+    --hash=sha256:d569388c381b24671589335a3be6e1d45546c2988c2ebe30fdcada8457a31008 \
+    --hash=sha256:d618649d4e70ac6efcbba75be98b26ef5078faad23592f9b51ca492953012429 \
+    --hash=sha256:d83a047959d38a7ff552ff94be767b7fd79b831ad1cd9920662db05fec24fe72 \
+    --hash=sha256:d8fff389528cad1618fb4b26b95550327495462cd745d879a8c7c2115248e399 \
+    --hash=sha256:da1758c76f50c39a2efd5e9859ce7d776317eb1dd34317c8152ac9251fc574a3 \
+    --hash=sha256:db7457bac39421addd0c8449933ac32d8042aae84a14911a757ae6ca3eef1392 \
+    --hash=sha256:e27bbb6d14416713a8bd7aaa1313c0fc8d44ee48d74497a0ff4c3a1b6ccb5167 \
+    --hash=sha256:e617fb6b0b6953fffd762669610c1c4ffd05632c138d61ac7e14ad187870669c \
+    --hash=sha256:e9aa71e15d9d9beaad2c6b9319edcdc0a49a43ef5c0a4c8265ca9ee7d6c67774 \
+    --hash=sha256:ec2abea24d98246b94913b76a125e855eb5c434f7c46546046372fe60f666351 \
+    --hash=sha256:f179dee3b863ab1c59580ff60f9d99f632f34ccb38bf67a33ec6b3ecadd0fd76 \
+    --hash=sha256:f4c035da3f544b1882bac24115f3e2e8760f10a0107614fc9839fd232200b875 \
+    --hash=sha256:f67f217af4b1ff66c68a87318012de788dd95fcfeb24cc889011f4e1c7454dfd \
+    --hash=sha256:f90c822a402cb865e396a504f9fc8173ef34212a342d92e362ca498cad308e28 \
+    --hash=sha256:ff3827aef427c89a25cc96ded1759271a93603aba9fb977a6d264648ebf989db
+    # via
+    #   aiohttp
+    #   yarl
 nest-asyncio==1.6.0 \
     --hash=sha256:6f172d5449aca15afd6c646851f4e31e02c598d553a667e38cafa997cfec55fe \
     --hash=sha256:87af6efd6b5e897c81050477ef65c62e2b2f35d51703cae01aff2905b1852e1c
     # via orbax-checkpoint
+neuralgcm==1.0.0 \
+    --hash=sha256:28f42b73188df92e35d5e187ef9ebd2b2a9922cfc46e05e0ae5202f451ef7c18 \
+    --hash=sha256:e75889894e8948211cec9bf18c89df5e223ef5a6d9a6444bb42f4b36143a5c27
+    # via -r builddeps/test-requirements.txt
+numcodecs==0.13.1 \
+    --hash=sha256:233bc7f26abce24d57e44ea8ebeb5cd17084690b4e7409dd470fdb75528d615f \
+    --hash=sha256:237b7171609e868a20fd313748494444458ccd696062f67e198f7f8f52000c15 \
+    --hash=sha256:2a86f5367af9168e30f99727ff03b27d849c31ad4522060dde0bce2923b3a8bc \
+    --hash=sha256:2eda97dd2f90add98df6d295f2c6ae846043396e3d51a739ca5db6c03b5eb666 \
+    --hash=sha256:3501a848adaddce98a71a262fee15cd3618312692aa419da77acd18af4a6a3f6 \
+    --hash=sha256:3f593c7506b0ab248961a3b13cb148cc6e8355662ff124ac591822310bc55ecf \
+    --hash=sha256:5195bea384a6428f8afcece793860b1ab0ae28143c853f0b2b20d55a8947c917 \
+    --hash=sha256:796b3e6740107e4fa624cc636248a1580138b3f1c579160f260f76ff13a4261b \
+    --hash=sha256:7a60d75179fd6692e301ddfb3b266d51eb598606dcae7b9fc57f986e8d65cb43 \
+    --hash=sha256:80d3071465f03522e776a31045ddf2cfee7f52df468b977ed3afdd7fe5869701 \
+    --hash=sha256:90d3065ae74c9342048ae0046006f99dcb1388b7288da5a19b3bddf9c30c3176 \
+    --hash=sha256:96add4f783c5ce57cc7e650b6cac79dd101daf887c479a00a29bc1487ced180b \
+    --hash=sha256:96e42f73c31b8c24259c5fac6adba0c3ebf95536e37749dc6c62ade2989dca28 \
+    --hash=sha256:a3cf37881df0898f3a9c0d4477df88133fe85185bffe57ba31bcc2fa207709bc \
+    --hash=sha256:da2230484e6102e5fa3cc1a5dd37ca1f92dfbd183d91662074d6f7574e3e8f53 \
+    --hash=sha256:e5db4824ebd5389ea30e54bc8aeccb82d514d28b6b68da6c536b8fa4596f4bca \
+    --hash=sha256:eda7d7823c9282e65234731fd6bd3986b1f9e035755f7fed248d7d366bb291ab
+    # via zarr
 numpy==2.1.0 \
     --hash=sha256:08801848a40aea24ce16c2ecde3b756f9ad756586fb2d13210939eb69b023f5b \
     --hash=sha256:0937e54c09f7a9a68da6889362ddd2ff584c02d015ec92672c099b61555f8911 \
@@ -426,6 +919,7 @@ numpy==2.1.0 \
     # via
     #   -r builddeps/test-requirements.txt
     #   chex
+    #   dinosaur-dycore
     #   dm-haiku
     #   e3nn-jax
     #   flax
@@ -435,11 +929,19 @@ numpy==2.1.0 \
     #   jmp
     #   jraph
     #   ml-dtypes
+    #   neuralgcm
+    #   numcodecs
     #   opt-einsum
     #   optax
     #   orbax-checkpoint
+    #   pandas
+    #   scikit-learn
     #   scipy
+    #   tensorflow-probability
     #   tensorstore
+    #   xarray
+    #   xarray-tensorstore
+    #   zarr
 nvidia-cublas-cu12==12.6.3.3 \
     --hash=sha256:e1f70bee38b964eac1907293b336bceb24498a4243e61eaf91a52977c59aebc4 \
     --hash=sha256:e531199ca4f1f764fb45bc1dde49a006f6765033f9c89c737e4553b9502ca1f5 \
@@ -498,6 +1000,10 @@ nvidia-nvjitlink-cu12==12.6.77 \
     #   nvidia-cufft-cu12
     #   nvidia-cusolver-cu12
     #   nvidia-cusparse-cu12
+oauthlib==3.2.2 \
+    --hash=sha256:8139f29aac13e25d502680e9e19963e83f16838d48a0d71c287fe40e7067fbca \
+    --hash=sha256:9859c40929662bec5d64f34d01c99e093149682a3f38915dc0655d5a633dd918
+    # via requests-oauthlib
 opt-einsum==3.3.0 \
     --hash=sha256:2455e59e3947d3c275477df7f5205b30635e266fe6dc300e3d9f9646bfcea147 \
     --hash=sha256:59f6475f77bbc37dcf7cd748519c0ec60722e91e63ca114e68821c0c54a46549
@@ -508,10 +1014,174 @@ optax==0.2.3 \
     # via
     #   flax
     #   jax-md
+    #   neuralgcm
 orbax-checkpoint==0.6.4 \
     --hash=sha256:366b4d528a7322e1b3d9ddcaed45c8515add0d2fc69c8975c30d98638543240f \
     --hash=sha256:b4f2608ee4d1da67f7619fc35ff9c928ecdf4ccf7546eeb43ecf38c2608b6dea
     # via flax
+packaging==24.1 \
+    --hash=sha256:026ed72c8ed3fcce5bf8950572258698927fd1dbda10a5e981cdf0ac37f4f002 \
+    --hash=sha256:5b8f2217dbdbd2f7f384c41c628544e6d52f2d0f53c6d0c3ea61aa5d1d7ff124
+    # via xarray
+pandas==2.2.3 \
+    --hash=sha256:062309c1b9ea12a50e8ce661145c6aab431b1e99530d3cd60640e255778bd43a \
+    --hash=sha256:15c0e1e02e93116177d29ff83e8b1619c93ddc9c49083f237d4312337a61165d \
+    --hash=sha256:1948ddde24197a0f7add2bdc4ca83bf2b1ef84a1bc8ccffd95eda17fd836ecb5 \
+    --hash=sha256:1db71525a1538b30142094edb9adc10be3f3e176748cd7acc2240c2f2e5aa3a4 \
+    --hash=sha256:22a9d949bfc9a502d320aa04e5d02feab689d61da4e7764b62c30b991c42c5f0 \
+    --hash=sha256:29401dbfa9ad77319367d36940cd8a0b3a11aba16063e39632d98b0e931ddf32 \
+    --hash=sha256:31d0ced62d4ea3e231a9f228366919a5ea0b07440d9d4dac345376fd8e1477ea \
+    --hash=sha256:3508d914817e153ad359d7e069d752cdd736a247c322d932eb89e6bc84217f28 \
+    --hash=sha256:37e0aced3e8f539eccf2e099f65cdb9c8aa85109b0be6e93e2baff94264bdc6f \
+    --hash=sha256:381175499d3802cde0eabbaf6324cce0c4f5d52ca6f8c377c29ad442f50f6348 \
+    --hash=sha256:38cf8125c40dae9d5acc10fa66af8ea6fdf760b2714ee482ca691fc66e6fcb18 \
+    --hash=sha256:3b71f27954685ee685317063bf13c7709a7ba74fc996b84fc6821c59b0f06468 \
+    --hash=sha256:3fc6873a41186404dad67245896a6e440baacc92f5b716ccd1bc9ed2995ab2c5 \
+    --hash=sha256:4850ba03528b6dd51d6c5d273c46f183f39a9baf3f0143e566b89450965b105e \
+    --hash=sha256:4f18ba62b61d7e192368b84517265a99b4d7ee8912f8708660fb4a366cc82667 \
+    --hash=sha256:56534ce0746a58afaf7942ba4863e0ef81c9c50d3f0ae93e9497d6a41a057645 \
+    --hash=sha256:59ef3764d0fe818125a5097d2ae867ca3fa64df032331b7e0917cf5d7bf66b13 \
+    --hash=sha256:5dbca4c1acd72e8eeef4753eeca07de9b1db4f398669d5994086f788a5d7cc30 \
+    --hash=sha256:5de54125a92bb4d1c051c0659e6fcb75256bf799a732a87184e5ea503965bce3 \
+    --hash=sha256:61c5ad4043f791b61dd4752191d9f07f0ae412515d59ba8f005832a532f8736d \
+    --hash=sha256:6374c452ff3ec675a8f46fd9ab25c4ad0ba590b71cf0656f8b6daa5202bca3fb \
+    --hash=sha256:63cc132e40a2e084cf01adf0775b15ac515ba905d7dcca47e9a251819c575ef3 \
+    --hash=sha256:66108071e1b935240e74525006034333f98bcdb87ea116de573a6a0dccb6c039 \
+    --hash=sha256:6dfcb5ee8d4d50c06a51c2fffa6cff6272098ad6540aed1a76d15fb9318194d8 \
+    --hash=sha256:7c2875855b0ff77b2a64a0365e24455d9990730d6431b9e0ee18ad8acee13dbd \
+    --hash=sha256:7eee9e7cea6adf3e3d24e304ac6b8300646e2a5d1cd3a3c2abed9101b0846761 \
+    --hash=sha256:800250ecdadb6d9c78eae4990da62743b857b470883fa27f652db8bdde7f6659 \
+    --hash=sha256:86976a1c5b25ae3f8ccae3a5306e443569ee3c3faf444dfd0f41cda24667ad57 \
+    --hash=sha256:8cd6d7cc958a3910f934ea8dbdf17b2364827bb4dafc38ce6eef6bb3d65ff09c \
+    --hash=sha256:99df71520d25fade9db7c1076ac94eb994f4d2673ef2aa2e86ee039b6746d20c \
+    --hash=sha256:a5a1595fe639f5988ba6a8e5bc9649af3baf26df3998a0abe56c02609392e0a4 \
+    --hash=sha256:ad5b65698ab28ed8d7f18790a0dc58005c7629f227be9ecc1072aa74c0c1d43a \
+    --hash=sha256:b1d432e8d08679a40e2a6d8b2f9770a5c21793a6f9f47fdd52c5ce1948a5a8a9 \
+    --hash=sha256:b8661b0238a69d7aafe156b7fa86c44b881387509653fdf857bebc5e4008ad42 \
+    --hash=sha256:ba96630bc17c875161df3818780af30e43be9b166ce51c9a18c1feae342906c2 \
+    --hash=sha256:bc6b93f9b966093cb0fd62ff1a7e4c09e6d546ad7c1de191767baffc57628f39 \
+    --hash=sha256:c124333816c3a9b03fbeef3a9f230ba9a737e9e5bb4060aa2107a86cc0a497fc \
+    --hash=sha256:cd8d0c3be0515c12fed0bdbae072551c8b54b7192c7b1fda0ba56059a0179698 \
+    --hash=sha256:d9c45366def9a3dd85a6454c0e7908f2b3b8e9c138f5dc38fed7ce720d8453ed \
+    --hash=sha256:f00d1345d84d8c86a63e476bb4955e46458b304b9575dcf71102b5c705320015 \
+    --hash=sha256:f3a255b2c19987fbbe62a9dfd6cff7ff2aa9ccab3fc75218fd4b7530f01efa24 \
+    --hash=sha256:fffb8ae78d8af97f849404f21411c95062db1496aeb3e56f146f0355c9989319
+    # via
+    #   dinosaur-dycore
+    #   neuralgcm
+    #   xarray
+pint==0.24.4 \
+    --hash=sha256:35275439b574837a6cd3020a5a4a73645eb125ce4152a73a2f126bf164b91b80 \
+    --hash=sha256:aa54926c8772159fcf65f82cc0d34de6768c151b32ad1deb0331291c38fe7659
+    # via dinosaur-dycore
+platformdirs==4.3.6 \
+    --hash=sha256:357fb2acbc885b0419afd3ce3ed34564c13c9b95c89360cd9563f73aa5e2b907 \
+    --hash=sha256:73e575e1408ab8103900836b97580d5307456908a03e92031bab39e4554cc3fb
+    # via pint
+propcache==0.2.0 \
+    --hash=sha256:00181262b17e517df2cd85656fcd6b4e70946fe62cd625b9d74ac9977b64d8d9 \
+    --hash=sha256:0e53cb83fdd61cbd67202735e6a6687a7b491c8742dfc39c9e01e80354956763 \
+    --hash=sha256:1235c01ddaa80da8235741e80815ce381c5267f96cc49b1477fdcf8c047ef325 \
+    --hash=sha256:140fbf08ab3588b3468932974a9331aff43c0ab8a2ec2c608b6d7d1756dbb6cb \
+    --hash=sha256:191db28dc6dcd29d1a3e063c3be0b40688ed76434622c53a284e5427565bbd9b \
+    --hash=sha256:1e41d67757ff4fbc8ef2af99b338bfb955010444b92929e9e55a6d4dcc3c4f09 \
+    --hash=sha256:1ec43d76b9677637a89d6ab86e1fef70d739217fefa208c65352ecf0282be957 \
+    --hash=sha256:20a617c776f520c3875cf4511e0d1db847a076d720714ae35ffe0df3e440be68 \
+    --hash=sha256:218db2a3c297a3768c11a34812e63b3ac1c3234c3a086def9c0fee50d35add1f \
+    --hash=sha256:22aa8f2272d81d9317ff5756bb108021a056805ce63dd3630e27d042c8092798 \
+    --hash=sha256:25a1f88b471b3bc911d18b935ecb7115dff3a192b6fef46f0bfaf71ff4f12418 \
+    --hash=sha256:25c8d773a62ce0451b020c7b29a35cfbc05de8b291163a7a0f3b7904f27253e6 \
+    --hash=sha256:2a60ad3e2553a74168d275a0ef35e8c0a965448ffbc3b300ab3a5bb9956c2162 \
+    --hash=sha256:2a66df3d4992bc1d725b9aa803e8c5a66c010c65c741ad901e260ece77f58d2f \
+    --hash=sha256:2ccc28197af5313706511fab3a8b66dcd6da067a1331372c82ea1cb74285e036 \
+    --hash=sha256:2e900bad2a8456d00a113cad8c13343f3b1f327534e3589acc2219729237a2e8 \
+    --hash=sha256:2ee7606193fb267be4b2e3b32714f2d58cad27217638db98a60f9efb5efeccc2 \
+    --hash=sha256:33ac8f098df0585c0b53009f039dfd913b38c1d2edafed0cedcc0c32a05aa110 \
+    --hash=sha256:3444cdba6628accf384e349014084b1cacd866fbb88433cd9d279d90a54e0b23 \
+    --hash=sha256:363ea8cd3c5cb6679f1c2f5f1f9669587361c062e4899fce56758efa928728f8 \
+    --hash=sha256:375a12d7556d462dc64d70475a9ee5982465fbb3d2b364f16b86ba9135793638 \
+    --hash=sha256:388f3217649d6d59292b722d940d4d2e1e6a7003259eb835724092a1cca0203a \
+    --hash=sha256:3947483a381259c06921612550867b37d22e1df6d6d7e8361264b6d037595f44 \
+    --hash=sha256:39e104da444a34830751715f45ef9fc537475ba21b7f1f5b0f4d71a3b60d7fe2 \
+    --hash=sha256:3c997f8c44ec9b9b0bcbf2d422cc00a1d9b9c681f56efa6ca149a941e5560da2 \
+    --hash=sha256:3dfafb44f7bb35c0c06eda6b2ab4bfd58f02729e7c4045e179f9a861b07c9850 \
+    --hash=sha256:3ebbcf2a07621f29638799828b8d8668c421bfb94c6cb04269130d8de4fb7136 \
+    --hash=sha256:3f88a4095e913f98988f5b338c1d4d5d07dbb0b6bad19892fd447484e483ba6b \
+    --hash=sha256:439e76255daa0f8151d3cb325f6dd4a3e93043e6403e6491813bcaaaa8733887 \
+    --hash=sha256:4569158070180c3855e9c0791c56be3ceeb192defa2cdf6a3f39e54319e56b89 \
+    --hash=sha256:466c219deee4536fbc83c08d09115249db301550625c7fef1c5563a584c9bc87 \
+    --hash=sha256:4a9d9b4d0a9b38d1c391bb4ad24aa65f306c6f01b512e10a8a34a2dc5675d348 \
+    --hash=sha256:4c7dde9e533c0a49d802b4f3f218fa9ad0a1ce21f2c2eb80d5216565202acab4 \
+    --hash=sha256:53d1bd3f979ed529f0805dd35ddaca330f80a9a6d90bc0121d2ff398f8ed8861 \
+    --hash=sha256:55346705687dbd7ef0d77883ab4f6fabc48232f587925bdaf95219bae072491e \
+    --hash=sha256:56295eb1e5f3aecd516d91b00cfd8bf3a13991de5a479df9e27dd569ea23959c \
+    --hash=sha256:56bb5c98f058a41bb58eead194b4db8c05b088c93d94d5161728515bd52b052b \
+    --hash=sha256:5a5b3bb545ead161be780ee85a2b54fdf7092815995661947812dde94a40f6fb \
+    --hash=sha256:5f2564ec89058ee7c7989a7b719115bdfe2a2fb8e7a4543b8d1c0cc4cf6478c1 \
+    --hash=sha256:608cce1da6f2672a56b24a015b42db4ac612ee709f3d29f27a00c943d9e851de \
+    --hash=sha256:63f13bf09cc3336eb04a837490b8f332e0db41da66995c9fd1ba04552e516354 \
+    --hash=sha256:662dd62358bdeaca0aee5761de8727cfd6861432e3bb828dc2a693aa0471a563 \
+    --hash=sha256:676135dcf3262c9c5081cc8f19ad55c8a64e3f7282a21266d05544450bffc3a5 \
+    --hash=sha256:67aeb72e0f482709991aa91345a831d0b707d16b0257e8ef88a2ad246a7280bf \
+    --hash=sha256:67b69535c870670c9f9b14a75d28baa32221d06f6b6fa6f77a0a13c5a7b0a5b9 \
+    --hash=sha256:682a7c79a2fbf40f5dbb1eb6bfe2cd865376deeac65acf9beb607505dced9e12 \
+    --hash=sha256:6994984550eaf25dd7fc7bd1b700ff45c894149341725bb4edc67f0ffa94efa4 \
+    --hash=sha256:69d3a98eebae99a420d4b28756c8ce6ea5a29291baf2dc9ff9414b42676f61d5 \
+    --hash=sha256:6e2e54267980349b723cff366d1e29b138b9a60fa376664a157a342689553f71 \
+    --hash=sha256:73e4b40ea0eda421b115248d7e79b59214411109a5bc47d0d48e4c73e3b8fcf9 \
+    --hash=sha256:74acd6e291f885678631b7ebc85d2d4aec458dd849b8c841b57ef04047833bed \
+    --hash=sha256:7665f04d0c7f26ff8bb534e1c65068409bf4687aa2534faf7104d7182debb336 \
+    --hash=sha256:7735e82e3498c27bcb2d17cb65d62c14f1100b71723b68362872bca7d0913d90 \
+    --hash=sha256:77a86c261679ea5f3896ec060be9dc8e365788248cc1e049632a1be682442063 \
+    --hash=sha256:7cf18abf9764746b9c8704774d8b06714bcb0a63641518a3a89c7f85cc02c2ad \
+    --hash=sha256:83928404adf8fb3d26793665633ea79b7361efa0287dfbd372a7e74311d51ee6 \
+    --hash=sha256:8e40876731f99b6f3c897b66b803c9e1c07a989b366c6b5b475fafd1f7ba3fb8 \
+    --hash=sha256:8f188cfcc64fb1266f4684206c9de0e80f54622c3f22a910cbd200478aeae61e \
+    --hash=sha256:91997d9cb4a325b60d4e3f20967f8eb08dfcb32b22554d5ef78e6fd1dda743a2 \
+    --hash=sha256:91ee8fc02ca52e24bcb77b234f22afc03288e1dafbb1f88fe24db308910c4ac7 \
+    --hash=sha256:92fe151145a990c22cbccf9ae15cae8ae9eddabfc949a219c9f667877e40853d \
+    --hash=sha256:945db8ee295d3af9dbdbb698cce9bbc5c59b5c3fe328bbc4387f59a8a35f998d \
+    --hash=sha256:9517d5e9e0731957468c29dbfd0f976736a0e55afaea843726e887f36fe017df \
+    --hash=sha256:952e0d9d07609d9c5be361f33b0d6d650cd2bae393aabb11d9b719364521984b \
+    --hash=sha256:97a58a28bcf63284e8b4d7b460cbee1edaab24634e82059c7b8c09e65284f178 \
+    --hash=sha256:97e48e8875e6c13909c800fa344cd54cc4b2b0db1d5f911f840458a500fde2c2 \
+    --hash=sha256:9e0f07b42d2a50c7dd2d8675d50f7343d998c64008f1da5fef888396b7f84630 \
+    --hash=sha256:a3dc1a4b165283bd865e8f8cb5f0c64c05001e0718ed06250d8cac9bec115b48 \
+    --hash=sha256:a3ebe9a75be7ab0b7da2464a77bb27febcb4fab46a34f9288f39d74833db7f61 \
+    --hash=sha256:a64e32f8bd94c105cc27f42d3b658902b5bcc947ece3c8fe7bc1b05982f60e89 \
+    --hash=sha256:a6ed8db0a556343d566a5c124ee483ae113acc9a557a807d439bcecc44e7dfbb \
+    --hash=sha256:ad9c9b99b05f163109466638bd30ada1722abb01bbb85c739c50b6dc11f92dc3 \
+    --hash=sha256:b33d7a286c0dc1a15f5fc864cc48ae92a846df287ceac2dd499926c3801054a6 \
+    --hash=sha256:bc092ba439d91df90aea38168e11f75c655880c12782facf5cf9c00f3d42b562 \
+    --hash=sha256:c436130cc779806bdf5d5fae0d848713105472b8566b75ff70048c47d3961c5b \
+    --hash=sha256:c5869b8fd70b81835a6f187c5fdbe67917a04d7e52b6e7cc4e5fe39d55c39d58 \
+    --hash=sha256:c5ecca8f9bab618340c8e848d340baf68bcd8ad90a8ecd7a4524a81c1764b3db \
+    --hash=sha256:cfac69017ef97db2438efb854edf24f5a29fd09a536ff3a992b75990720cdc99 \
+    --hash=sha256:d2f0d0f976985f85dfb5f3d685697ef769faa6b71993b46b295cdbbd6be8cc37 \
+    --hash=sha256:d5bed7f9805cc29c780f3aee05de3262ee7ce1f47083cfe9f77471e9d6777e83 \
+    --hash=sha256:d6a21ef516d36909931a2967621eecb256018aeb11fc48656e3257e73e2e247a \
+    --hash=sha256:d9b6ddac6408194e934002a69bcaadbc88c10b5f38fb9307779d1c629181815d \
+    --hash=sha256:db47514ffdbd91ccdc7e6f8407aac4ee94cc871b15b577c1c324236b013ddd04 \
+    --hash=sha256:df81779732feb9d01e5d513fad0122efb3d53bbc75f61b2a4f29a020bc985e70 \
+    --hash=sha256:e4a91d44379f45f5e540971d41e4626dacd7f01004826a18cb048e7da7e96544 \
+    --hash=sha256:e63e3e1e0271f374ed489ff5ee73d4b6e7c60710e1f76af5f0e1a6117cd26394 \
+    --hash=sha256:e70fac33e8b4ac63dfc4c956fd7d85a0b1139adcfc0d964ce288b7c527537fea \
+    --hash=sha256:ecddc221a077a8132cf7c747d5352a15ed763b674c0448d811f408bf803d9ad7 \
+    --hash=sha256:f45eec587dafd4b2d41ac189c2156461ebd0c1082d2fe7013571598abb8505d1 \
+    --hash=sha256:f52a68c21363c45297aca15561812d542f8fc683c85201df0bebe209e349f793 \
+    --hash=sha256:f571aea50ba5623c308aa146eb650eebf7dbe0fd8c5d946e28343cb3b5aad577 \
+    --hash=sha256:f60f0ac7005b9f5a6091009b09a419ace1610e163fa5deaba5ce3484341840e7 \
+    --hash=sha256:f6475a1b2ecb310c98c28d271a30df74f9dd436ee46d09236a6b750a7599ce57 \
+    --hash=sha256:f6d5749fdd33d90e34c2efb174c7e236829147a2713334d708746e94c4bde40d \
+    --hash=sha256:f902804113e032e2cdf8c71015651c97af6418363bea8d78dc0911d56c335032 \
+    --hash=sha256:fa1076244f54bb76e65e22cb6910365779d5c3d71d1f18b275f1dfc7b0d71b4d \
+    --hash=sha256:fc2db02409338bf36590aa985a461b2c96fce91f8e7e0f14c50c5fcc4f229016 \
+    --hash=sha256:ffcad6c564fe6b9b8916c1aefbb37a362deebf9394bd2974e9d84232e3e08504
+    # via yarl
+proto-plus==1.25.0 \
+    --hash=sha256:c91fc4a65074ade8e458e95ef8bac34d4008daa7cce4a12d6707066fca648961 \
+    --hash=sha256:fbb17f57f7bd05a68b7707e745e26528b0b3c34e378db91eef93912c54982d91
+    # via google-api-core
 protobuf==5.28.2 \
     --hash=sha256:2c69461a7fcc8e24be697624c09a839976d82ae75062b11a0972e41fd2cd9132 \
     --hash=sha256:35cfcb15f213449af7ff6198d6eb5f739c37d7e4f1c09b5d0641babf2cc0c68f \
@@ -524,11 +1194,33 @@ protobuf==5.28.2 \
     --hash=sha256:c0ea0123dac3399a2eeb1a1443d82b7afc9ff40241433296769f7da42d142ec3 \
     --hash=sha256:ca53faf29896c526863366a52a8f4d88e69cd04ec9571ed6082fa117fac3ab36 \
     --hash=sha256:eeea10f3dc0ac7e6b4933d32db20662902b4ab81bf28df12218aa389e9c2102d
-    # via orbax-checkpoint
+    # via
+    #   google-api-core
+    #   googleapis-common-protos
+    #   orbax-checkpoint
+    #   proto-plus
+pyasn1==0.6.1 \
+    --hash=sha256:0d632f46f2ba09143da3a8afe9e33fb6f92fa2320ab7e886e2d0f7672af84629 \
+    --hash=sha256:6f580d2bdd84365380830acf45550f2511469f673cb4a5ae3857a3170128b034
+    # via
+    #   pyasn1-modules
+    #   rsa
+pyasn1-modules==0.4.1 \
+    --hash=sha256:49bfa96b45a292b711e986f222502c1c9a5e1f4e568fc30e2574a6c7d07838fd \
+    --hash=sha256:c28e2dbf9c06ad61c71a075c7e0f9fd0f1b0bb2d2ad4377f240d33ac2ab60a7c
+    # via google-auth
 pygments==2.18.0 \
     --hash=sha256:786ff802f32e91311bff3889f6e9a86e81505fe99f2735bb6d60ae0c5004f199 \
     --hash=sha256:b8e6aca0523f3ab76fee51799c488e38782ac06eafcf95e7ba832985c8e7b13a
     # via rich
+python-dateutil==2.9.0.post0 \
+    --hash=sha256:37dd54208da7e1cd875388217d5e00ebd4179249f90fb72437e91a35459a0ad3 \
+    --hash=sha256:a8b2bc7bffae282281c8140a97d3aa9c14da0b136dfe83f850eea9a5f7470427
+    # via pandas
+pytz==2024.2 \
+    --hash=sha256:2aa355083c50a0f93fa581709deac0c9ad65cca8a9e9beac660adcbd493c798a \
+    --hash=sha256:31c7c1817eb7fae7ca4b8c7ee50c72f93aa2dd863de768e1ef4245d426aa0725
+    # via pandas
 pyyaml==6.0.2 \
     --hash=sha256:01179a4a8559ab5de078078f37e5c1a30d76bb88519906844fd7bdea1b7729ff \
     --hash=sha256:0833f8694549e586547b576dcfaba4a6b55b9e96098b36cdc7ebefe667dfed48 \
@@ -590,11 +1282,52 @@ pyyaml==6.0.2 \
 requests==2.32.3 ; sys_platform == "linux" \
     --hash=sha256:55365417734eb18255590a9ff9eb97e9e1da868d4ccd6402399eaf68af20a760 \
     --hash=sha256:70761cfe03c773ceb22aa2f671b4757976145175cdfca038c02654d061d6dcc6
-    # via -r builddeps/test-requirements.txt
+    # via
+    #   -r builddeps/test-requirements.txt
+    #   gcsfs
+    #   google-api-core
+    #   google-cloud-storage
+    #   requests-oauthlib
+requests-oauthlib==2.0.0 \
+    --hash=sha256:7dd8a5c40426b779b0868c404bdef9768deccf22749cde15852df527e6269b36 \
+    --hash=sha256:b3dffaebd884d8cd778494369603a9e7b58d29111bf6b41bdc2dcd87203af4e9
+    # via google-auth-oauthlib
 rich==13.9.2 \
     --hash=sha256:51a2c62057461aaf7152b4d611168f93a9fc73068f8ded2790f29fe2b5366d0c \
     --hash=sha256:8c82a3d3f8dcfe9e734771313e606b39d8247bb6b826e196f4914b333b743cf1
     # via flax
+rsa==4.9 \
+    --hash=sha256:90260d9058e514786967344d0ef75fa8727eed8a7d2e43ce9f4bcf1b536174f7 \
+    --hash=sha256:e38464a49c6c85d7f1351b0126661487a7e0a14a50f1675ec50eb34d4f20ef21
+    # via google-auth
+scikit-learn==1.5.2 \
+    --hash=sha256:03b6158efa3faaf1feea3faa884c840ebd61b6484167c711548fce208ea09445 \
+    --hash=sha256:178ddd0a5cb0044464fc1bfc4cca5b1833bfc7bb022d70b05db8530da4bb3dd3 \
+    --hash=sha256:1ff45e26928d3b4eb767a8f14a9a6efbf1cbff7c05d1fb0f95f211a89fd4f5de \
+    --hash=sha256:299406827fb9a4f862626d0fe6c122f5f87f8910b86fe5daa4c32dcd742139b6 \
+    --hash=sha256:2d4cad1119c77930b235579ad0dc25e65c917e756fe80cab96aa3b9428bd3fb0 \
+    --hash=sha256:394397841449853c2290a32050382edaec3da89e35b3e03d6cc966aebc6a8ae6 \
+    --hash=sha256:3a686885a4b3818d9e62904d91b57fa757fc2bed3e465c8b177be652f4dd37c8 \
+    --hash=sha256:3b923d119d65b7bd555c73be5423bf06c0105678ce7e1f558cb4b40b0a5502b1 \
+    --hash=sha256:3bed4909ba187aca80580fe2ef370d9180dcf18e621a27c4cf2ef10d279a7efe \
+    --hash=sha256:52788f48b5d8bca5c0736c175fa6bdaab2ef00a8f536cda698db61bd89c551c1 \
+    --hash=sha256:57cc1786cfd6bd118220a92ede80270132aa353647684efa385a74244a41e3b1 \
+    --hash=sha256:643964678f4b5fbdc95cbf8aec638acc7aa70f5f79ee2cdad1eec3df4ba6ead8 \
+    --hash=sha256:6c16d84a0d45e4894832b3c4d0bf73050939e21b99b01b6fd59cbb0cf39163b6 \
+    --hash=sha256:757c7d514ddb00ae249832fe87100d9c73c6ea91423802872d9e74970a0e40b9 \
+    --hash=sha256:8c412ccc2ad9bf3755915e3908e677b367ebc8d010acbb3f182814524f2e5540 \
+    --hash=sha256:b0768ad641981f5d3a198430a1d31c3e044ed2e8a6f22166b4d546a5116d7908 \
+    --hash=sha256:b4237ed7b3fdd0a4882792e68ef2545d5baa50aca3bb45aa7df468138ad8f94d \
+    --hash=sha256:b7b0f9a0b1040830d38c39b91b3a44e1b643f4b36e36567b80b7c6bd2202a27f \
+    --hash=sha256:c15b1ca23d7c5f33cc2cb0a0d6aaacf893792271cddff0edbd6a40e8319bc113 \
+    --hash=sha256:ca64b3089a6d9b9363cd3546f8978229dcbb737aceb2c12144ee3f70f95684b7 \
+    --hash=sha256:e9a702e2de732bbb20d3bad29ebd77fc05a6b427dc49964300340e4c9328b3f5 \
+    --hash=sha256:f60021ec1574e56632be2a36b946f8143bf4e5e6af4a06d85281adc22938e0dd \
+    --hash=sha256:f7284ade780084d94505632241bf78c44ab3b6f1e8ccab3d2af58e0e950f9c12 \
+    --hash=sha256:f763897fe92d0e903aa4847b0aec0e68cadfff77e8a0687cabd946c89d17e675 \
+    --hash=sha256:f8b0ccd4a902836493e026c03256e8b206656f91fbcc4fde28c57a5b752561f1 \
+    --hash=sha256:f932a02c3f4956dfb981391ab24bda1dbd90fe3d628e4b42caef3e041c67707a
+    # via dinosaur-dycore
 scipy==1.14.1 \
     --hash=sha256:0c2f95de3b04e26f5f3ad5bb05e74ba7f68b837133a4492414b3afd79dfe540e \
     --hash=sha256:1729560c906963fc8389f6aac023739ff3983e727b1a4d87696b7bf108316a79 \
@@ -630,12 +1363,17 @@ scipy==1.14.1 \
     --hash=sha256:edaf02b82cd7639db00dbff629995ef185c8df4c3ffa71a5562a595765a06ce1 \
     --hash=sha256:fef8c87f8abfb884dac04e97824b61299880c43f4ce675dd2cbeadd3c9b466d2
     # via
+    #   dinosaur-dycore
     #   jax
     #   jaxlib
+    #   scikit-learn
 six==1.16.0 \
     --hash=sha256:1e61c37477a1626458e36f7b1d82aa5c9b094fa4802892072e49de9c60c4c926 \
     --hash=sha256:8abb2f1d86890a2dfb989f9a77cfcfd3e47c2a354b01111771326f8aa26e0254
-    # via ml-collections
+    # via
+    #   ml-collections
+    #   python-dateutil
+    #   tensorflow-probability
 sympy==1.13.3 \
     --hash=sha256:54612cf55a62755ee71824ce692986f23c88ffa77207b30c1368eda4a7060f73 \
     --hash=sha256:b27fd2c6530e0ab39e275fc9b683895367e51d5da91baa8d3d64db2565fec4d9
@@ -644,6 +1382,9 @@ tabulate==0.9.0 \
     --hash=sha256:0095b12bf5966de529c0feb1fa08671671b3368eec77d7ef7ab114be2c068b3c \
     --hash=sha256:024ca478df22e9340661486f85298cff5f6dcdba14f3813e8830015b9ed1948f
     # via dm-haiku
+tensorflow-probability==0.24.0 \
+    --hash=sha256:8c1774683e38359dbcaf3697e79b7e6a4e69b9c7b3679e78ee18f43e59e5759b
+    # via neuralgcm
 tensorstore==0.1.66 \
     --hash=sha256:0261b1b7121186ac4898c9b4e29289c46754f9d7d73e906620e9a21386fe564e \
     --hash=sha256:06f7dba630d75f65201b1c021fcd89e5f8538506b444b5dd712f35c11b812fc8 \
@@ -669,10 +1410,19 @@ tensorstore==0.1.66 \
     # via
     #   flax
     #   orbax-checkpoint
+    #   xarray-tensorstore
+threadpoolctl==3.5.0 \
+    --hash=sha256:082433502dd922bf738de0d8bcc4fdcbf0979ff44c42bd40f5af8a282f6fa107 \
+    --hash=sha256:56c1e26c150397e58c4926da8eeee87533b1e32bef131bd4bf6a2f45f3185467
+    # via scikit-learn
 toolz==1.0.0 \
     --hash=sha256:292c8f1c4e7516bf9086f8850935c799a874039c8bcf959d47b600e4c44a6236 \
     --hash=sha256:2c86e3d9a04798ac556793bced838816296a2f085017664e4995cb40a1047a02
     # via chex
+tree-math==0.2.1 \
+    --hash=sha256:77613c2360e067b16c0738b32e7d64aa5e1167e609ca1a32b6fa05dcb0d330bb \
+    --hash=sha256:80b01a5d29c1339db0cd8b465f9fad4f872adac117ed599b14f093a6c52821fe
+    # via dinosaur-dycore
 typing-extensions==4.12.2 \
     --hash=sha256:04e5ca0351e0f3f85c6853954072df659d0d13fac324d0072316b67d7794700d \
     --hash=sha256:1a7ead55c7e559dd4dee8856e3a88b41225abfe1ce8df57b7c13915fe121ffb8
@@ -680,11 +1430,117 @@ typing-extensions==4.12.2 \
     #   chex
     #   etils
     #   flax
+    #   flexcache
+    #   flexparser
     #   orbax-checkpoint
+    #   pint
+tzdata==2024.2 \
+    --hash=sha256:7d85cc416e9382e69095b7bdf4afd9e3880418a2413feec7069d533d6b4e31cc \
+    --hash=sha256:a48093786cdcde33cad18c2555e8532f34422074448fbc874186f0abd79565cd
+    # via pandas
 urllib3==2.2.2 \
     --hash=sha256:a448b2f64d686155468037e1ace9f2d2199776e17f0a46610480d311f73e3472 \
     --hash=sha256:dd505485549a7a552833da5e6063639d0d177c04f23bc3864e41e5dc5f612168
     # via requests
+xarray==2024.10.0 \
+    --hash=sha256:ae1d38cb44a0324dfb61e492394158ae22389bf7de9f3c174309c17376df63a0 \
+    --hash=sha256:e369e2bac430e418c2448e5b96f07da4635f98c1319aa23cfeb3fbcb9a01d2e0
+    # via
+    #   dinosaur-dycore
+    #   neuralgcm
+    #   xarray-tensorstore
+xarray-tensorstore==0.1.4 \
+    --hash=sha256:1dacec6376912373dbf9b239d96184090f400a49cd66017f5bf0c824793c9552 \
+    --hash=sha256:de3558c08891c4a7524cf579cb271d6cc005962efd1b8601dd679b58b13a68c8
+    # via dinosaur-dycore
+yarl==1.17.1 \
+    --hash=sha256:06157fb3c58f2736a5e47c8fcbe1afc8b5de6fb28b14d25574af9e62150fcaac \
+    --hash=sha256:067a63fcfda82da6b198fa73079b1ca40b7c9b7994995b6ee38acda728b64d47 \
+    --hash=sha256:0b1794853124e2f663f0ea54efb0340b457f08d40a1cef78edfa086576179c91 \
+    --hash=sha256:0bdff5e0995522706c53078f531fb586f56de9c4c81c243865dd5c66c132c3b5 \
+    --hash=sha256:117ed8b3732528a1e41af3aa6d4e08483c2f0f2e3d3d7dca7cf538b3516d93df \
+    --hash=sha256:14bc88baa44e1f84164a392827b5defb4fa8e56b93fecac3d15315e7c8e5d8b3 \
+    --hash=sha256:1654ec814b18be1af2c857aa9000de7a601400bd4c9ca24629b18486c2e35463 \
+    --hash=sha256:16bca6678a83657dd48df84b51bd56a6c6bd401853aef6d09dc2506a78484c7b \
+    --hash=sha256:1a3b91c44efa29e6c8ef8a9a2b583347998e2ba52c5d8280dbd5919c02dfc3b5 \
+    --hash=sha256:1a52a1ffdd824fb1835272e125385c32fd8b17fbdefeedcb4d543cc23b332d74 \
+    --hash=sha256:1ce36ded585f45b1e9bb36d0ae94765c6608b43bd2e7f5f88079f7a85c61a4d3 \
+    --hash=sha256:299f11b44d8d3a588234adbe01112126010bd96d9139c3ba7b3badd9829261c3 \
+    --hash=sha256:2b24ec55fad43e476905eceaf14f41f6478780b870eda5d08b4d6de9a60b65b4 \
+    --hash=sha256:2d374d70fdc36f5863b84e54775452f68639bc862918602d028f89310a034ab0 \
+    --hash=sha256:2d9f0606baaec5dd54cb99667fcf85183a7477f3766fbddbe3f385e7fc253299 \
+    --hash=sha256:2e7ba4c9377e48fb7b20dedbd473cbcbc13e72e1826917c185157a137dac9df2 \
+    --hash=sha256:2f0a6423295a0d282d00e8701fe763eeefba8037e984ad5de44aa349002562ac \
+    --hash=sha256:327828786da2006085a4d1feb2594de6f6d26f8af48b81eb1ae950c788d97f61 \
+    --hash=sha256:380e6c38ef692b8fd5a0f6d1fa8774d81ebc08cfbd624b1bca62a4d4af2f9931 \
+    --hash=sha256:3b74ff4767d3ef47ffe0cd1d89379dc4d828d4873e5528976ced3b44fe5b0a21 \
+    --hash=sha256:3e844be8d536afa129366d9af76ed7cb8dfefec99f5f1c9e4f8ae542279a6dc3 \
+    --hash=sha256:459e81c2fb920b5f5df744262d1498ec2c8081acdcfe18181da44c50f51312f7 \
+    --hash=sha256:46ddf6e0b975cd680eb83318aa1d321cb2bf8d288d50f1754526230fcf59ba96 \
+    --hash=sha256:482c122b72e3c5ec98f11457aeb436ae4aecca75de19b3d1de7cf88bc40db82f \
+    --hash=sha256:561c87fea99545ef7d692403c110b2f99dced6dff93056d6e04384ad3bc46243 \
+    --hash=sha256:578d00c9b7fccfa1745a44f4eddfdc99d723d157dad26764538fbdda37209857 \
+    --hash=sha256:58c8e9620eb82a189c6c40cb6b59b4e35b2ee68b1f2afa6597732a2b467d7e8f \
+    --hash=sha256:5b29beab10211a746f9846baa39275e80034e065460d99eb51e45c9a9495bcca \
+    --hash=sha256:5d1d42556b063d579cae59e37a38c61f4402b47d70c29f0ef15cee1acaa64488 \
+    --hash=sha256:5f236cb5999ccd23a0ab1bd219cfe0ee3e1c1b65aaf6dd3320e972f7ec3a39da \
+    --hash=sha256:62a91aefff3d11bf60e5956d340eb507a983a7ec802b19072bb989ce120cd948 \
+    --hash=sha256:64cc6e97f14cf8a275d79c5002281f3040c12e2e4220623b5759ea7f9868d6a5 \
+    --hash=sha256:6f4c9156c4d1eb490fe374fb294deeb7bc7eaccda50e23775b2354b6a6739934 \
+    --hash=sha256:7294e38f9aa2e9f05f765b28ffdc5d81378508ce6dadbe93f6d464a8c9594473 \
+    --hash=sha256:7615058aabad54416ddac99ade09a5510cf77039a3b903e94e8922f25ed203d7 \
+    --hash=sha256:7e48cdb8226644e2fbd0bdb0a0f87906a3db07087f4de77a1b1b1ccfd9e93685 \
+    --hash=sha256:7f63d176a81555984e91f2c84c2a574a61cab7111cc907e176f0f01538e9ff6e \
+    --hash=sha256:7f6595c852ca544aaeeb32d357e62c9c780eac69dcd34e40cae7b55bc4fb1147 \
+    --hash=sha256:7fac95714b09da9278a0b52e492466f773cfe37651cf467a83a1b659be24bf71 \
+    --hash=sha256:81713b70bea5c1386dc2f32a8f0dab4148a2928c7495c808c541ee0aae614d67 \
+    --hash=sha256:846dd2e1243407133d3195d2d7e4ceefcaa5f5bf7278f0a9bda00967e6326b04 \
+    --hash=sha256:84c063af19ef5130084db70ada40ce63a84f6c1ef4d3dbc34e5e8c4febb20822 \
+    --hash=sha256:881764d610e3269964fc4bb3c19bb6fce55422828e152b885609ec176b41cf11 \
+    --hash=sha256:8994b29c462de9a8fce2d591028b986dbbe1b32f3ad600b2d3e1c482c93abad6 \
+    --hash=sha256:8c79e9d7e3d8a32d4824250a9c6401194fb4c2ad9a0cec8f6a96e09a582c2cc0 \
+    --hash=sha256:8ee427208c675f1b6e344a1f89376a9613fc30b52646a04ac0c1f6587c7e46ec \
+    --hash=sha256:949681f68e0e3c25377462be4b658500e85ca24323d9619fdc41f68d46a1ffda \
+    --hash=sha256:9e275792097c9f7e80741c36de3b61917aebecc08a67ae62899b074566ff8556 \
+    --hash=sha256:9fb815155aac6bfa8d86184079652c9715c812d506b22cfa369196ef4e99d1b4 \
+    --hash=sha256:a2a64e62c7a0edd07c1c917b0586655f3362d2c2d37d474db1a509efb96fea1c \
+    --hash=sha256:a7ac5b4984c468ce4f4a553df281450df0a34aefae02e58d77a0847be8d1e11f \
+    --hash=sha256:aa46dce75078fceaf7cecac5817422febb4355fbdda440db55206e3bd288cfb8 \
+    --hash=sha256:ae3476e934b9d714aa8000d2e4c01eb2590eee10b9d8cd03e7983ad65dfbfcba \
+    --hash=sha256:b0341e6d9a0c0e3cdc65857ef518bb05b410dbd70d749a0d33ac0f39e81a4258 \
+    --hash=sha256:b40d1bf6e6f74f7c0a567a9e5e778bbd4699d1d3d2c0fe46f4b717eef9e96b95 \
+    --hash=sha256:b5c4804e4039f487e942c13381e6c27b4b4e66066d94ef1fae3f6ba8b953f383 \
+    --hash=sha256:b5d6a6c9602fd4598fa07e0389e19fe199ae96449008d8304bf5d47cb745462e \
+    --hash=sha256:b5f1ac7359e17efe0b6e5fec21de34145caef22b260e978336f325d5c84e6938 \
+    --hash=sha256:c0167540094838ee9093ef6cc2c69d0074bbf84a432b4995835e8e5a0d984374 \
+    --hash=sha256:c180ac742a083e109c1a18151f4dd8675f32679985a1c750d2ff806796165b55 \
+    --hash=sha256:c73df5b6e8fabe2ddb74876fb82d9dd44cbace0ca12e8861ce9155ad3c886139 \
+    --hash=sha256:c7e177c619342e407415d4f35dec63d2d134d951e24b5166afcdfd1362828e17 \
+    --hash=sha256:cbad927ea8ed814622305d842c93412cb47bd39a496ed0f96bfd42b922b4a217 \
+    --hash=sha256:cc353841428d56b683a123a813e6a686e07026d6b1c5757970a877195f880c2d \
+    --hash=sha256:cc7c92c1baa629cb03ecb0c3d12564f172218fb1739f54bf5f3881844daadc6d \
+    --hash=sha256:cc7d768260f4ba4ea01741c1b5fe3d3a6c70eb91c87f4c8761bbcce5181beafe \
+    --hash=sha256:d0eea830b591dbc68e030c86a9569826145df485b2b4554874b07fea1275a199 \
+    --hash=sha256:d216e5d9b8749563c7f2c6f7a0831057ec844c68b4c11cb10fc62d4fd373c26d \
+    --hash=sha256:d401f07261dc5aa36c2e4efc308548f6ae943bfff20fcadb0a07517a26b196d8 \
+    --hash=sha256:d6324274b4e0e2fa1b3eccb25997b1c9ed134ff61d296448ab8269f5ac068c4c \
+    --hash=sha256:d8a8b74d843c2638f3864a17d97a4acda58e40d3e44b6303b8cc3d3c44ae2d29 \
+    --hash=sha256:d9b6b28a57feb51605d6ae5e61a9044a31742db557a3b851a74c13bc61de5172 \
+    --hash=sha256:de599af166970d6a61accde358ec9ded821234cbbc8c6413acfec06056b8e860 \
+    --hash=sha256:e594b22688d5747b06e957f1ef822060cb5cb35b493066e33ceac0cf882188b7 \
+    --hash=sha256:e5b078134f48552c4d9527db2f7da0b5359abd49393cdf9794017baec7506170 \
+    --hash=sha256:eb6dce402734575e1a8cc0bb1509afca508a400a57ce13d306ea2c663bad1138 \
+    --hash=sha256:f1790a4b1e8e8e028c391175433b9c8122c39b46e1663228158e61e6f915bf06 \
+    --hash=sha256:f5efe0661b9fcd6246f27957f6ae1c0eb29bc60552820f01e970b4996e016004 \
+    --hash=sha256:f9cbfbc5faca235fbdf531b93aa0f9f005ec7d267d9d738761a4d42b744ea159 \
+    --hash=sha256:fbea1751729afe607d84acfd01efd95e3b31db148a181a441984ce9b3d3469da \
+    --hash=sha256:fca4b4307ebe9c3ec77a084da3a9d1999d164693d16492ca2b64594340999988 \
+    --hash=sha256:ff5c6771c7e3511a06555afa317879b7db8d640137ba55d6ab0d0c50425cab75
+    # via aiohttp
+zarr==2.18.3 \
+    --hash=sha256:2580d8cb6dd84621771a10d31c4d777dca8a27706a1a89b29f42d2d37e2df5ce \
+    --hash=sha256:b1f7dfd2496f436745cdd4c7bcf8d3b4bc1dceef5fdd0d589c87130d842496dd
+    # via xarray-tensorstore
 zipp==3.20.2 \
     --hash=sha256:a817ac80d6cf4b23bf7f2828b7cabf326f15a001bea8b1f9b49631780ba28350 \
     --hash=sha256:bc9eb26f4506fda01b81bcde0ca78103b6e62f991b381fec825435c836edbc29
diff --git a/test/BUILD b/test/BUILD
index 6f514a24..6e83bbfc 100644
--- a/test/BUILD
+++ b/test/BUILD
@@ -137,9 +137,9 @@ py_test(
 )
 
 py_test(
-    name = "neuralgcm",
+    name = "neuralgcm_test",
     srcs = [
-        "neuralgcm.py",
+        "neuralgcm_test.py",
         "test_utils.py",
     ],
     imports = ["."],

From be4c81dc7967f53d119e8aad658ea0544a4b6d6a Mon Sep 17 00:00:00 2001
From: William Moses <gh@wsmoses.com>
Date: Thu, 7 Nov 2024 17:04:16 -0600
Subject: [PATCH 14/17] Update jaxmd.py

---
 test/jaxmd.py | 6 ------
 1 file changed, 6 deletions(-)

diff --git a/test/jaxmd.py b/test/jaxmd.py
index 40243544..700995ea 100644
--- a/test/jaxmd.py
+++ b/test/jaxmd.py
@@ -24,12 +24,6 @@
 
 jax.config.update("jax_enable_x64", True)
 
-from jax_md import space
-from jax_md import energy
-from jax_md import simulate
-from jax_md import quantity
-from jax_md import partition
-
 pipelines = [
     ("JaX  ", None, CurBackends),
     ("JaXPipe", JaXPipeline(), CurBackends),

From 9627c06ebbc0e8f94e78c1b745485d643b2d63f0 Mon Sep 17 00:00:00 2001
From: William Moses <gh@wsmoses.com>
Date: Thu, 7 Nov 2024 17:06:07 -0600
Subject: [PATCH 15/17] Update neuralgcm_test.py

---
 test/neuralgcm_test.py | 14 ++++++++++----
 1 file changed, 10 insertions(+), 4 deletions(-)

diff --git a/test/neuralgcm_test.py b/test/neuralgcm_test.py
index 9db4aac1..a9f3b2e9 100644
--- a/test/neuralgcm_test.py
+++ b/test/neuralgcm_test.py
@@ -93,10 +93,16 @@ def setUp(self):
         eval_era5 = xarray_utils.regrid(sliced_era5, regridder)
         eval_era5 = xarray_utils.fill_nan_with_nearest(eval_era5)
 
-        # inner_steps = 24  # save model outputs once every 24 hours
-        # outer_steps = 4 * 24 // inner_steps  # total of 4 days
-        inner_steps = 4  # save model outputs once every 24 hours
-        outer_steps = 4 * 4 // inner_steps  # total of 4 days
+        import os
+        if os.getenv("NEURALGCM_LARGE") is not None: 
+            inner_steps = 24  # save model outputs once every 24 hours
+            outer_steps = 4 * 24 // inner_steps  # total of 4 days
+        elif s.getenv("NEURALGCM_MEDIUM") is not None: 
+            inner_steps = 4  # save model outputs once every 24 hours
+            outer_steps = 4 * 4 // inner_steps  # total of 4 days
+        else:
+            inner_steps = 1  # save model outputs once every 24 hours
+            outer_steps = 1 * 1 // inner_steps  # total of 4 days
 
         timedelta = np.timedelta64(1, "h") * inner_steps
         # times = (np.arange(outer_steps) * inner_steps)  # time axis in hours

From 4de383184d76e2e38c4f00d663a31f2ee612e437 Mon Sep 17 00:00:00 2001
From: "William S. Moses" <gh@wsmoses.com>
Date: Thu, 7 Nov 2024 18:07:01 -0500
Subject: [PATCH 16/17] fix

---
 test/neuralgcm_test.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/test/neuralgcm_test.py b/test/neuralgcm_test.py
index a9f3b2e9..a919f7d7 100644
--- a/test/neuralgcm_test.py
+++ b/test/neuralgcm_test.py
@@ -94,10 +94,11 @@ def setUp(self):
         eval_era5 = xarray_utils.fill_nan_with_nearest(eval_era5)
 
         import os
-        if os.getenv("NEURALGCM_LARGE") is not None: 
+
+        if os.getenv("NEURALGCM_LARGE") is not None:
             inner_steps = 24  # save model outputs once every 24 hours
             outer_steps = 4 * 24 // inner_steps  # total of 4 days
-        elif s.getenv("NEURALGCM_MEDIUM") is not None: 
+        elif s.getenv("NEURALGCM_MEDIUM") is not None:
             inner_steps = 4  # save model outputs once every 24 hours
             outer_steps = 4 * 4 // inner_steps  # total of 4 days
         else:

From 925ee41c9cac4534272d75900d8b5d8a78c69f87 Mon Sep 17 00:00:00 2001
From: William Moses <gh@wsmoses.com>
Date: Thu, 7 Nov 2024 18:03:06 -0600
Subject: [PATCH 17/17] Update neuralgcm_test.py

---
 test/neuralgcm_test.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/test/neuralgcm_test.py b/test/neuralgcm_test.py
index a919f7d7..47051a19 100644
--- a/test/neuralgcm_test.py
+++ b/test/neuralgcm_test.py
@@ -102,8 +102,8 @@ def setUp(self):
             inner_steps = 4  # save model outputs once every 24 hours
             outer_steps = 4 * 4 // inner_steps  # total of 4 days
         else:
-            inner_steps = 1  # save model outputs once every 24 hours
-            outer_steps = 1 * 1 // inner_steps  # total of 4 days
+            inner_steps = 2  # save model outputs once every 24 hours
+            outer_steps = 2 * 2 // inner_steps  # total of 4 days
 
         timedelta = np.timedelta64(1, "h") * inner_steps
         # times = (np.arange(outer_steps) * inner_steps)  # time axis in hours