--fanout FLAG added to [weatherbench2]/scripts/compute_zonal_energy_s…

…pectrum PiperOrigin-RevId: 575340755
google-research · Oct 20, 2023 · 131e05e · 131e05e
1 parent 58d6f14
commit 131e05e
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/scripts/compute_zonal_energy_spectrum.py b/scripts/compute_zonal_energy_spectrum.py
@@ -91,6 +91,11 @@
         ' average.'
     ),
 )
+FANOUT = flags.DEFINE_integer(
+    'fanout',
+    None,
+    help='Beam CombineFn fanout. Might be required for large dataset.',
+)
 
 RUNNER = flags.DEFINE_string('runner', None, 'beam.runners.Runner')
 
@@ -199,7 +204,7 @@ def main(argv: list[str]) -> None:
         )
         | xbeam.SplitVariables()
         | beam.MapTuple(_strip_offsets)
-        | xbeam.Mean(AVERAGING_DIMS.value)
+        | xbeam.Mean(AVERAGING_DIMS.value, fanout=FANOUT.value)
         | xbeam.ChunksToZarr(
             OUTPUT_PATH.value, template, output_chunks, num_threads=16
         )