From 2ad99db4e5fc4046adaed3827b9b05d3540919c0 Mon Sep 17 00:00:00 2001 From: Jake VanderPlas Date: Mon, 16 Jun 2025 08:31:31 -0700 Subject: [PATCH] Silence some pytype errors related to a JAX build refactor This build change allows pytype to propagate annotations that it previously did not, and because of this it starts flagging existing incorrect annotations. PiperOrigin-RevId: 772038309 --- clu/__init__.py | 2 +- clu/asynclib.py | 2 +- clu/asynclib_test.py | 2 +- clu/checkpoint.py | 2 +- clu/checkpoint_test.py | 2 +- clu/data/__init__.py | 2 +- clu/data/dataset_iterator.py | 2 +- clu/data/dataset_iterator_test.py | 2 +- clu/deterministic_data.py | 2 +- clu/deterministic_data_test.py | 2 +- clu/internal/__init__.py | 2 +- clu/internal/utils.py | 2 +- clu/internal/utils_test.py | 2 +- clu/metric_writers/__init__.py | 2 +- clu/metric_writers/async_writer.py | 2 +- clu/metric_writers/async_writer_test.py | 2 +- clu/metric_writers/interface.py | 2 +- clu/metric_writers/logging_writer.py | 2 +- clu/metric_writers/logging_writer_test.py | 2 +- clu/metric_writers/multi_writer.py | 2 +- clu/metric_writers/multi_writer_test.py | 2 +- clu/metric_writers/summary_writer.py | 2 +- clu/metric_writers/tf/__init__.py | 2 +- clu/metric_writers/tf/summary_writer.py | 2 +- clu/metric_writers/tf/summary_writer_test.py | 2 +- clu/metric_writers/torch_tensorboard_writer.py | 2 +- clu/metric_writers/torch_tensorboard_writer_test.py | 2 +- clu/metric_writers/utils.py | 2 +- clu/metric_writers/utils_test.py | 2 +- clu/metrics.py | 4 ++-- clu/metrics_test.py | 2 +- clu/parameter_overview.py | 2 +- clu/parameter_overview_test.py | 2 +- clu/periodic_actions.py | 2 +- clu/periodic_actions_test.py | 2 +- clu/platform/__init__.py | 2 +- clu/platform/interface.py | 2 +- clu/platform/local.py | 2 +- clu/preprocess_spec.py | 2 +- clu/preprocess_spec_test.py | 2 +- clu/profiler.py | 2 +- clu/values.py | 2 +- setup.py | 2 +- 43 files changed, 44 insertions(+), 44 deletions(-) diff --git a/clu/__init__.py b/clu/__init__.py index bf7fdc7..5ec1547 100644 --- a/clu/__init__.py +++ b/clu/__init__.py @@ -1,4 +1,4 @@ -# Copyright 2024 The CLU Authors. +# Copyright 2025 The CLU Authors. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/clu/asynclib.py b/clu/asynclib.py index 05f4e3c..e63fa8f 100644 --- a/clu/asynclib.py +++ b/clu/asynclib.py @@ -1,4 +1,4 @@ -# Copyright 2024 The CLU Authors. +# Copyright 2025 The CLU Authors. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/clu/asynclib_test.py b/clu/asynclib_test.py index 38a359c..ce36c5c 100644 --- a/clu/asynclib_test.py +++ b/clu/asynclib_test.py @@ -1,4 +1,4 @@ -# Copyright 2024 The CLU Authors. +# Copyright 2025 The CLU Authors. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/clu/checkpoint.py b/clu/checkpoint.py index 1560685..9b086b6 100644 --- a/clu/checkpoint.py +++ b/clu/checkpoint.py @@ -1,4 +1,4 @@ -# Copyright 2024 The CLU Authors. +# Copyright 2025 The CLU Authors. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/clu/checkpoint_test.py b/clu/checkpoint_test.py index 78ffe68..a42744f 100644 --- a/clu/checkpoint_test.py +++ b/clu/checkpoint_test.py @@ -1,4 +1,4 @@ -# Copyright 2024 The CLU Authors. +# Copyright 2025 The CLU Authors. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/clu/data/__init__.py b/clu/data/__init__.py index 16eeeb5..ff94e80 100644 --- a/clu/data/__init__.py +++ b/clu/data/__init__.py @@ -1,4 +1,4 @@ -# Copyright 2024 The CLU Authors. +# Copyright 2025 The CLU Authors. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/clu/data/dataset_iterator.py b/clu/data/dataset_iterator.py index 2fe9a85..3254d06 100644 --- a/clu/data/dataset_iterator.py +++ b/clu/data/dataset_iterator.py @@ -1,4 +1,4 @@ -# Copyright 2024 The CLU Authors. +# Copyright 2025 The CLU Authors. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/clu/data/dataset_iterator_test.py b/clu/data/dataset_iterator_test.py index 3472123..6485811 100644 --- a/clu/data/dataset_iterator_test.py +++ b/clu/data/dataset_iterator_test.py @@ -1,4 +1,4 @@ -# Copyright 2024 The CLU Authors. +# Copyright 2025 The CLU Authors. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/clu/deterministic_data.py b/clu/deterministic_data.py index 6fb4671..5d88160 100644 --- a/clu/deterministic_data.py +++ b/clu/deterministic_data.py @@ -1,4 +1,4 @@ -# Copyright 2024 The CLU Authors. +# Copyright 2025 The CLU Authors. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/clu/deterministic_data_test.py b/clu/deterministic_data_test.py index 2460fb6..aef0251 100644 --- a/clu/deterministic_data_test.py +++ b/clu/deterministic_data_test.py @@ -1,4 +1,4 @@ -# Copyright 2024 The CLU Authors. +# Copyright 2025 The CLU Authors. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/clu/internal/__init__.py b/clu/internal/__init__.py index bf7fdc7..5ec1547 100644 --- a/clu/internal/__init__.py +++ b/clu/internal/__init__.py @@ -1,4 +1,4 @@ -# Copyright 2024 The CLU Authors. +# Copyright 2025 The CLU Authors. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/clu/internal/utils.py b/clu/internal/utils.py index fea71e9..9238597 100644 --- a/clu/internal/utils.py +++ b/clu/internal/utils.py @@ -1,4 +1,4 @@ -# Copyright 2024 The CLU Authors. +# Copyright 2025 The CLU Authors. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/clu/internal/utils_test.py b/clu/internal/utils_test.py index ffb129f..6da370d 100644 --- a/clu/internal/utils_test.py +++ b/clu/internal/utils_test.py @@ -1,4 +1,4 @@ -# Copyright 2024 The CLU Authors. +# Copyright 2025 The CLU Authors. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/clu/metric_writers/__init__.py b/clu/metric_writers/__init__.py index e67cde7..c7604b0 100644 --- a/clu/metric_writers/__init__.py +++ b/clu/metric_writers/__init__.py @@ -1,4 +1,4 @@ -# Copyright 2024 The CLU Authors. +# Copyright 2025 The CLU Authors. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/clu/metric_writers/async_writer.py b/clu/metric_writers/async_writer.py index 06ea5ff..21fa1ff 100644 --- a/clu/metric_writers/async_writer.py +++ b/clu/metric_writers/async_writer.py @@ -1,4 +1,4 @@ -# Copyright 2024 The CLU Authors. +# Copyright 2025 The CLU Authors. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/clu/metric_writers/async_writer_test.py b/clu/metric_writers/async_writer_test.py index c3d10a5..dd844da 100644 --- a/clu/metric_writers/async_writer_test.py +++ b/clu/metric_writers/async_writer_test.py @@ -1,4 +1,4 @@ -# Copyright 2024 The CLU Authors. +# Copyright 2025 The CLU Authors. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/clu/metric_writers/interface.py b/clu/metric_writers/interface.py index e14988d..9cec258 100644 --- a/clu/metric_writers/interface.py +++ b/clu/metric_writers/interface.py @@ -1,4 +1,4 @@ -# Copyright 2024 The CLU Authors. +# Copyright 2025 The CLU Authors. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/clu/metric_writers/logging_writer.py b/clu/metric_writers/logging_writer.py index 5cddb56..e608936 100644 --- a/clu/metric_writers/logging_writer.py +++ b/clu/metric_writers/logging_writer.py @@ -1,4 +1,4 @@ -# Copyright 2024 The CLU Authors. +# Copyright 2025 The CLU Authors. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/clu/metric_writers/logging_writer_test.py b/clu/metric_writers/logging_writer_test.py index 977fff4..c39bdae 100644 --- a/clu/metric_writers/logging_writer_test.py +++ b/clu/metric_writers/logging_writer_test.py @@ -1,4 +1,4 @@ -# Copyright 2024 The CLU Authors. +# Copyright 2025 The CLU Authors. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/clu/metric_writers/multi_writer.py b/clu/metric_writers/multi_writer.py index f1f0827..86649b2 100644 --- a/clu/metric_writers/multi_writer.py +++ b/clu/metric_writers/multi_writer.py @@ -1,4 +1,4 @@ -# Copyright 2024 The CLU Authors. +# Copyright 2025 The CLU Authors. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/clu/metric_writers/multi_writer_test.py b/clu/metric_writers/multi_writer_test.py index 350ea9a..ef63258 100644 --- a/clu/metric_writers/multi_writer_test.py +++ b/clu/metric_writers/multi_writer_test.py @@ -1,4 +1,4 @@ -# Copyright 2024 The CLU Authors. +# Copyright 2025 The CLU Authors. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/clu/metric_writers/summary_writer.py b/clu/metric_writers/summary_writer.py index c9e495a..db2b77c 100644 --- a/clu/metric_writers/summary_writer.py +++ b/clu/metric_writers/summary_writer.py @@ -1,4 +1,4 @@ -# Copyright 2024 The CLU Authors. +# Copyright 2025 The CLU Authors. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/clu/metric_writers/tf/__init__.py b/clu/metric_writers/tf/__init__.py index fbb6d34..fd4b9da 100644 --- a/clu/metric_writers/tf/__init__.py +++ b/clu/metric_writers/tf/__init__.py @@ -1,4 +1,4 @@ -# Copyright 2024 The CLU Authors. +# Copyright 2025 The CLU Authors. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/clu/metric_writers/tf/summary_writer.py b/clu/metric_writers/tf/summary_writer.py index 7e00570..0c3e852 100644 --- a/clu/metric_writers/tf/summary_writer.py +++ b/clu/metric_writers/tf/summary_writer.py @@ -1,4 +1,4 @@ -# Copyright 2024 The CLU Authors. +# Copyright 2025 The CLU Authors. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/clu/metric_writers/tf/summary_writer_test.py b/clu/metric_writers/tf/summary_writer_test.py index a52990e..8ac2932 100644 --- a/clu/metric_writers/tf/summary_writer_test.py +++ b/clu/metric_writers/tf/summary_writer_test.py @@ -1,4 +1,4 @@ -# Copyright 2024 The CLU Authors. +# Copyright 2025 The CLU Authors. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/clu/metric_writers/torch_tensorboard_writer.py b/clu/metric_writers/torch_tensorboard_writer.py index 2d413ca..50bc945 100644 --- a/clu/metric_writers/torch_tensorboard_writer.py +++ b/clu/metric_writers/torch_tensorboard_writer.py @@ -1,4 +1,4 @@ -# Copyright 2024 The CLU Authors. +# Copyright 2025 The CLU Authors. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/clu/metric_writers/torch_tensorboard_writer_test.py b/clu/metric_writers/torch_tensorboard_writer_test.py index 449fbaf..a6290fe 100644 --- a/clu/metric_writers/torch_tensorboard_writer_test.py +++ b/clu/metric_writers/torch_tensorboard_writer_test.py @@ -1,4 +1,4 @@ -# Copyright 2024 The CLU Authors. +# Copyright 2025 The CLU Authors. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/clu/metric_writers/utils.py b/clu/metric_writers/utils.py index a98b310..a43d964 100644 --- a/clu/metric_writers/utils.py +++ b/clu/metric_writers/utils.py @@ -1,4 +1,4 @@ -# Copyright 2024 The CLU Authors. +# Copyright 2025 The CLU Authors. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/clu/metric_writers/utils_test.py b/clu/metric_writers/utils_test.py index a004470..cad7301 100644 --- a/clu/metric_writers/utils_test.py +++ b/clu/metric_writers/utils_test.py @@ -1,4 +1,4 @@ -# Copyright 2024 The CLU Authors. +# Copyright 2025 The CLU Authors. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/clu/metrics.py b/clu/metrics.py index d8e247d..803f2ae 100644 --- a/clu/metrics.py +++ b/clu/metrics.py @@ -1,4 +1,4 @@ -# Copyright 2024 The CLU Authors. +# Copyright 2025 The CLU Authors. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -198,7 +198,7 @@ def reduce_step(reduced: M, metric: M) -> tuple[M, None]: # According to b/160868467#comment4, usage of `jax.lax.scan` does not add a # significant computational cost for simple metrics where e.g. `jnp.sum` # could be used instead. - return jax.lax.scan(reduce_step, first, remainder)[0] + return jax.lax.scan(reduce_step, first, remainder)[0] # pytype: disable=wrong-arg-types # lax-types @classmethod def from_fun(cls, fun: FromFunCallable): # No way to annotate return type diff --git a/clu/metrics_test.py b/clu/metrics_test.py index baf389e..dd4c5fc 100644 --- a/clu/metrics_test.py +++ b/clu/metrics_test.py @@ -1,4 +1,4 @@ -# Copyright 2024 The CLU Authors. +# Copyright 2025 The CLU Authors. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/clu/parameter_overview.py b/clu/parameter_overview.py index 8cdb921..c45ccd9 100644 --- a/clu/parameter_overview.py +++ b/clu/parameter_overview.py @@ -1,4 +1,4 @@ -# Copyright 2024 The CLU Authors. +# Copyright 2025 The CLU Authors. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/clu/parameter_overview_test.py b/clu/parameter_overview_test.py index 9fdf0a6..c8a2d13 100644 --- a/clu/parameter_overview_test.py +++ b/clu/parameter_overview_test.py @@ -1,4 +1,4 @@ -# Copyright 2024 The CLU Authors. +# Copyright 2025 The CLU Authors. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/clu/periodic_actions.py b/clu/periodic_actions.py index 7fe87a7..75ba08b 100644 --- a/clu/periodic_actions.py +++ b/clu/periodic_actions.py @@ -1,4 +1,4 @@ -# Copyright 2024 The CLU Authors. +# Copyright 2025 The CLU Authors. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/clu/periodic_actions_test.py b/clu/periodic_actions_test.py index bf70e97..d31a7eb 100644 --- a/clu/periodic_actions_test.py +++ b/clu/periodic_actions_test.py @@ -1,4 +1,4 @@ -# Copyright 2024 The CLU Authors. +# Copyright 2025 The CLU Authors. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/clu/platform/__init__.py b/clu/platform/__init__.py index fd1fd0d..7ae5eef 100644 --- a/clu/platform/__init__.py +++ b/clu/platform/__init__.py @@ -1,4 +1,4 @@ -# Copyright 2024 The CLU Authors. +# Copyright 2025 The CLU Authors. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/clu/platform/interface.py b/clu/platform/interface.py index 8c8e0db..21ccd4a 100644 --- a/clu/platform/interface.py +++ b/clu/platform/interface.py @@ -1,4 +1,4 @@ -# Copyright 2024 The CLU Authors. +# Copyright 2025 The CLU Authors. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/clu/platform/local.py b/clu/platform/local.py index 51f0b37..534051a 100644 --- a/clu/platform/local.py +++ b/clu/platform/local.py @@ -1,4 +1,4 @@ -# Copyright 2024 The CLU Authors. +# Copyright 2025 The CLU Authors. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/clu/preprocess_spec.py b/clu/preprocess_spec.py index 83d18a8..fba8602 100644 --- a/clu/preprocess_spec.py +++ b/clu/preprocess_spec.py @@ -1,4 +1,4 @@ -# Copyright 2024 The CLU Authors. +# Copyright 2025 The CLU Authors. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/clu/preprocess_spec_test.py b/clu/preprocess_spec_test.py index 6a32235..8e9f85f 100644 --- a/clu/preprocess_spec_test.py +++ b/clu/preprocess_spec_test.py @@ -1,4 +1,4 @@ -# Copyright 2024 The CLU Authors. +# Copyright 2025 The CLU Authors. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/clu/profiler.py b/clu/profiler.py index f946f8a..11e6b4c 100644 --- a/clu/profiler.py +++ b/clu/profiler.py @@ -1,4 +1,4 @@ -# Copyright 2024 The CLU Authors. +# Copyright 2025 The CLU Authors. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/clu/values.py b/clu/values.py index a3204ee..a6da91d 100644 --- a/clu/values.py +++ b/clu/values.py @@ -1,4 +1,4 @@ -# Copyright 2024 The CLU Authors. +# Copyright 2025 The CLU Authors. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/setup.py b/setup.py index a403e87..d6761b8 100644 --- a/setup.py +++ b/setup.py @@ -1,4 +1,4 @@ -# Copyright 2024 The CLU Authors. +# Copyright 2025 The CLU Authors. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License.