From ea2515bb135ff7331a10c3aeb4e7da9f9821c7d7 Mon Sep 17 00:00:00 2001 From: Ian Langmore Date: Mon, 4 Dec 2023 11:25:37 -0800 Subject: [PATCH] Pytype fix in WB2/compute_statistical_moments.py PiperOrigin-RevId: 587787044 --- scripts/compute_statistical_moments.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/scripts/compute_statistical_moments.py b/scripts/compute_statistical_moments.py index 94c6ee8..31ab6bb 100644 --- a/scripts/compute_statistical_moments.py +++ b/scripts/compute_statistical_moments.py @@ -154,7 +154,8 @@ def main(argv: list[str]) -> None: ) # Rechunk in time pcoll_time = pcoll_tmp | f'RechunkTime_{order}' >> xbeam.Rechunk( - space_reduce_template.sizes, + # Convert to string to satisfy pytype. + {str(k): v for k, v in space_reduce_template.sizes.items()}, reduce_working_chunks, time_working_chunks, itemsize=RECHUNK_ITEMSIZE.value,