google · copybara-service · Feb 26, 2026
diff --git a/clu/__init__.py b/clu/__init__.py
@@ -1,4 +1,4 @@
-# Copyright 2025 The CLU Authors.
+# Copyright 2026 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
@@ -1,4 +1,4 @@
-# Copyright 2025 The CLU Authors.
+# Copyright 2026 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
@@ -1,4 +1,4 @@
-# Copyright 2025 The CLU Authors.
+# Copyright 2026 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
@@ -1,4 +1,4 @@
-# Copyright 2025 The CLU Authors.
+# Copyright 2026 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
@@ -1,4 +1,4 @@
-# Copyright 2025 The CLU Authors.
+# Copyright 2026 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
@@ -1,4 +1,4 @@
-# Copyright 2025 The CLU Authors.
+# Copyright 2026 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
@@ -1,4 +1,4 @@
-# Copyright 2025 The CLU Authors.
+# Copyright 2026 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
@@ -1,4 +1,4 @@
-# Copyright 2025 The CLU Authors.
+# Copyright 2026 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
@@ -1,4 +1,4 @@
-# Copyright 2025 The CLU Authors.
+# Copyright 2026 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
@@ -1,4 +1,4 @@
-# Copyright 2025 The CLU Authors.
+# Copyright 2026 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
@@ -1,4 +1,4 @@
-# Copyright 2025 The CLU Authors.
+# Copyright 2026 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
@@ -1,4 +1,4 @@
-# Copyright 2025 The CLU Authors.
+# Copyright 2026 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
@@ -1,4 +1,4 @@
-# Copyright 2025 The CLU Authors.
+# Copyright 2026 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
@@ -1,4 +1,4 @@
-# Copyright 2025 The CLU Authors.
+# Copyright 2026 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
@@ -1,4 +1,4 @@
-# Copyright 2025 The CLU Authors.
+# Copyright 2026 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
@@ -1,4 +1,4 @@
-# Copyright 2025 The CLU Authors.
+# Copyright 2026 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
@@ -1,4 +1,4 @@
-# Copyright 2025 The CLU Authors.
+# Copyright 2026 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
@@ -1,4 +1,4 @@
-# Copyright 2025 The CLU Authors.
+# Copyright 2026 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
@@ -1,4 +1,4 @@
-# Copyright 2025 The CLU Authors.
+# Copyright 2026 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
@@ -1,4 +1,4 @@
-# Copyright 2025 The CLU Authors.
+# Copyright 2026 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
@@ -1,4 +1,4 @@
-# Copyright 2025 The CLU Authors.
+# Copyright 2026 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
@@ -1,4 +1,4 @@
-# Copyright 2025 The CLU Authors.
+# Copyright 2026 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
@@ -1,4 +1,4 @@
-# Copyright 2025 The CLU Authors.
+# Copyright 2026 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
@@ -1,4 +1,4 @@
-# Copyright 2025 The CLU Authors.
+# Copyright 2026 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
@@ -1,4 +1,4 @@
-# Copyright 2025 The CLU Authors.
+# Copyright 2026 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
@@ -1,4 +1,4 @@
-# Copyright 2025 The CLU Authors.
+# Copyright 2026 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
@@ -1,4 +1,4 @@
-# Copyright 2025 The CLU Authors.
+# Copyright 2026 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
@@ -1,4 +1,4 @@
-# Copyright 2025 The CLU Authors.
+# Copyright 2026 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
@@ -1,4 +1,4 @@
-# Copyright 2025 The CLU Authors.
+# Copyright 2026 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
@@ -1,4 +1,4 @@
-# Copyright 2025 The CLU Authors.
+# Copyright 2026 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.
@@ -193,8 +193,7 @@ def reduce_step(reduced: M, metric: M) -> tuple[M, None]:
       # pylint: disable-next=protected-access
       return reduced._reduce_merge(metric), None
 
-    # Avoid degraded performance under the new jax.pmap. See
-    # https://docs.jax.dev/en/latest/migrate_pmap.html#int-indexing-into-sharded-arrays.
+    # Avoid degraded performance under the new jax.pmap.
     # Only use the sharding path for concrete sharded arrays, not tracers.
     def _is_concrete_sharded(x):
       if isinstance(x, jax.core.Tracer):
@@ -208,11 +207,7 @@ def _is_concrete_sharded(x):
       return shards[0].data.ndim > 0 and shards[0].data.shape[0] == 1
 
     leaves = jax.tree_util.tree_leaves(self)
-    use_sharding_path = (
-        jax.config.jax_pmap_shmap_merge
-        and leaves
-        and _is_concrete_sharded(leaves[0])
-    )
+    use_sharding_path = leaves and _is_concrete_sharded(leaves[0])
 
     if use_sharding_path:
 

diff --git a/clu/metrics_test.py b/clu/metrics_test.py
@@ -1,4 +1,4 @@
-# Copyright 2025 The CLU Authors.
+# Copyright 2026 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
@@ -1,4 +1,4 @@
-# Copyright 2025 The CLU Authors.
+# Copyright 2026 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
@@ -1,4 +1,4 @@
-# Copyright 2025 The CLU Authors.
+# Copyright 2026 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
@@ -1,4 +1,4 @@
-# Copyright 2025 The CLU Authors.
+# Copyright 2026 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
@@ -1,4 +1,4 @@
-# Copyright 2025 The CLU Authors.
+# Copyright 2026 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
@@ -1,4 +1,4 @@
-# Copyright 2025 The CLU Authors.
+# Copyright 2026 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
@@ -1,4 +1,4 @@
-# Copyright 2025 The CLU Authors.
+# Copyright 2026 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
@@ -1,4 +1,4 @@
-# Copyright 2025 The CLU Authors.
+# Copyright 2026 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
@@ -1,4 +1,4 @@
-# Copyright 2025 The CLU Authors.
+# Copyright 2026 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
@@ -1,4 +1,4 @@
-# Copyright 2025 The CLU Authors.
+# Copyright 2026 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
@@ -1,4 +1,4 @@
-# Copyright 2025 The CLU Authors.
+# Copyright 2026 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
@@ -1,4 +1,4 @@
-# Copyright 2025 The CLU Authors.
+# Copyright 2026 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
@@ -1,4 +1,4 @@
-# Copyright 2025 The CLU Authors.
+# Copyright 2026 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.