From d1dda49d32fc9712c8d4f2470031520f885e9a3c 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 --- precondition/__init__.py | 2 +- precondition/datamix_gemma/bandit_loop.py | 2 +- .../datamix_gemma/confusion_matrix_calc.py | 2 +- precondition/datamix_gemma/cross_compile.py | 2 +- .../datamix_gemma/dataset_builders/__init__.py | 2 +- .../dataset_builders/dataset_builder.py | 2 +- .../dataset_builders/gsm8k_dataset_builder.py | 2 +- .../dataset_builders/mbpp_dataset_builder.py | 2 +- .../dataset_builders/mtnt_dataset_builder.py | 2 +- .../dataset_builders/open_orca_dataset_builder.py | 2 +- .../dataset_builders/orca_math_dataset_builder.py | 2 +- .../preprocessed_codealpaca_dataset_builder.py | 2 +- .../preprocessed_dolly_dataset_builder.py | 2 +- .../preprocessed_gsm8k_dataset_builder.py | 2 +- .../preprocessed_metamath_dataset_builder.py | 2 +- .../preprocessed_open_orca_dataset_builder.py | 2 +- .../preprocessed_orca_math_dataset_builder.py | 2 +- .../preprocessed_sciq_dataset_builder.py | 2 +- .../preprocessed_wikipedia_dataset_builder.py | 2 +- .../datamix_gemma/deconstructed_sampler.py | 2 +- .../deterministic_strategy_bandit_loop.py | 2 +- precondition/datamix_gemma/evals/__init__.py | 2 +- precondition/datamix_gemma/evals/crop.py | 2 +- precondition/datamix_gemma/evals/eval.py | 2 +- precondition/datamix_gemma/evals/gsm8k_eval.py | 2 +- precondition/datamix_gemma/evals/mbpp_eval.py | 2 +- precondition/datamix_gemma/evals/mmlu_eval.py | 2 +- precondition/datamix_gemma/finetune.py | 2 +- .../datamix_gemma/finetune_eval_measurement.py | 2 +- precondition/datamix_gemma/finetune_utils.py | 2 +- .../datamix_gemma/finetuning_experiment.py | 2 +- precondition/datamix_gemma/random_baseline.py | 2 +- precondition/datamix_gemma/snr_calculation.py | 2 +- precondition/datamix_gemma/tokenizers/__init__.py | 2 +- .../datamix_gemma/tokenizers/gemma_tokenizer.py | 2 +- .../training_batch_generators/__init__.py | 2 +- ...board_deterministic_training_batch_generator.py | 2 +- ...mportance_weighting_training_batch_generator.py | 2 +- ...mportance_weighting_training_batch_generator.py | 2 +- ...mportance_weighting_training_batch_generator.py | 2 +- .../training_batch_generator.py | 2 +- .../vanilla_training_batch_generator.py | 2 +- precondition/datamix_gemma/training_loop.py | 2 +- precondition/distributed_shampoo.py | 14 +++++++------- precondition/distributed_shampoo_test.py | 2 +- precondition/oco/algorithms.py | 2 +- precondition/oco/datasets.py | 2 +- precondition/oco/sweep.py | 2 +- precondition/oco/train.py | 2 +- precondition/quantization_utils.py | 2 +- precondition/sm3.py | 2 +- precondition/sm3_test.py | 2 +- precondition/tearfree/grafting.py | 2 +- precondition/tearfree/grafting_test.py | 2 +- precondition/tearfree/momentum.py | 2 +- precondition/tearfree/momentum_test.py | 2 +- precondition/tearfree/optimizer.py | 2 +- precondition/tearfree/optimizer_smoke_test.py | 2 +- precondition/tearfree/optimizer_test.py | 2 +- precondition/tearfree/praxis_shim.py | 2 +- precondition/tearfree/reallocation.py | 2 +- precondition/tearfree/reallocation_test.py | 2 +- precondition/tearfree/reshaper.py | 2 +- precondition/tearfree/reshaper_test.py | 2 +- precondition/tearfree/second_order.py | 2 +- precondition/tearfree/shampoo.py | 2 +- precondition/tearfree/shampoo_test.py | 2 +- precondition/tearfree/sketchy.py | 2 +- precondition/tearfree/sketchy_test.py | 2 +- 69 files changed, 75 insertions(+), 75 deletions(-) diff --git a/precondition/__init__.py b/precondition/__init__.py index 01c8f63..2f00f3c 100644 --- a/precondition/__init__.py +++ b/precondition/__init__.py @@ -1,4 +1,4 @@ -# Copyright 2024 The precondition Authors. +# Copyright 2025 The precondition 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/precondition/datamix_gemma/bandit_loop.py b/precondition/datamix_gemma/bandit_loop.py index 9b01e23..271163b 100644 --- a/precondition/datamix_gemma/bandit_loop.py +++ b/precondition/datamix_gemma/bandit_loop.py @@ -1,4 +1,4 @@ -# Copyright 2024 The precondition Authors. +# Copyright 2025 The precondition 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/precondition/datamix_gemma/confusion_matrix_calc.py b/precondition/datamix_gemma/confusion_matrix_calc.py index 71e15fd..2345d3d 100644 --- a/precondition/datamix_gemma/confusion_matrix_calc.py +++ b/precondition/datamix_gemma/confusion_matrix_calc.py @@ -1,4 +1,4 @@ -# Copyright 2024 The precondition Authors. +# Copyright 2025 The precondition 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/precondition/datamix_gemma/cross_compile.py b/precondition/datamix_gemma/cross_compile.py index b46fd22..396ffff 100644 --- a/precondition/datamix_gemma/cross_compile.py +++ b/precondition/datamix_gemma/cross_compile.py @@ -1,4 +1,4 @@ -# Copyright 2024 The precondition Authors. +# Copyright 2025 The precondition 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/precondition/datamix_gemma/dataset_builders/__init__.py b/precondition/datamix_gemma/dataset_builders/__init__.py index 44d891e..7ee5dd2 100644 --- a/precondition/datamix_gemma/dataset_builders/__init__.py +++ b/precondition/datamix_gemma/dataset_builders/__init__.py @@ -1,4 +1,4 @@ -# Copyright 2024 The precondition Authors. +# Copyright 2025 The precondition 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/precondition/datamix_gemma/dataset_builders/dataset_builder.py b/precondition/datamix_gemma/dataset_builders/dataset_builder.py index 5dd09b2..8f4224b 100644 --- a/precondition/datamix_gemma/dataset_builders/dataset_builder.py +++ b/precondition/datamix_gemma/dataset_builders/dataset_builder.py @@ -1,4 +1,4 @@ -# Copyright 2024 The precondition Authors. +# Copyright 2025 The precondition 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/precondition/datamix_gemma/dataset_builders/gsm8k_dataset_builder.py b/precondition/datamix_gemma/dataset_builders/gsm8k_dataset_builder.py index 0fd1c50..acb2f26 100644 --- a/precondition/datamix_gemma/dataset_builders/gsm8k_dataset_builder.py +++ b/precondition/datamix_gemma/dataset_builders/gsm8k_dataset_builder.py @@ -1,4 +1,4 @@ -# Copyright 2024 The precondition Authors. +# Copyright 2025 The precondition 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/precondition/datamix_gemma/dataset_builders/mbpp_dataset_builder.py b/precondition/datamix_gemma/dataset_builders/mbpp_dataset_builder.py index 42ab424..b13e845 100644 --- a/precondition/datamix_gemma/dataset_builders/mbpp_dataset_builder.py +++ b/precondition/datamix_gemma/dataset_builders/mbpp_dataset_builder.py @@ -1,4 +1,4 @@ -# Copyright 2024 The precondition Authors. +# Copyright 2025 The precondition 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/precondition/datamix_gemma/dataset_builders/mtnt_dataset_builder.py b/precondition/datamix_gemma/dataset_builders/mtnt_dataset_builder.py index 709306e..e3c2d87 100644 --- a/precondition/datamix_gemma/dataset_builders/mtnt_dataset_builder.py +++ b/precondition/datamix_gemma/dataset_builders/mtnt_dataset_builder.py @@ -1,4 +1,4 @@ -# Copyright 2024 The precondition Authors. +# Copyright 2025 The precondition 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/precondition/datamix_gemma/dataset_builders/open_orca_dataset_builder.py b/precondition/datamix_gemma/dataset_builders/open_orca_dataset_builder.py index e2e0de7..ac94005 100644 --- a/precondition/datamix_gemma/dataset_builders/open_orca_dataset_builder.py +++ b/precondition/datamix_gemma/dataset_builders/open_orca_dataset_builder.py @@ -1,4 +1,4 @@ -# Copyright 2024 The precondition Authors. +# Copyright 2025 The precondition 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/precondition/datamix_gemma/dataset_builders/orca_math_dataset_builder.py b/precondition/datamix_gemma/dataset_builders/orca_math_dataset_builder.py index e0d83aa..245bac0 100644 --- a/precondition/datamix_gemma/dataset_builders/orca_math_dataset_builder.py +++ b/precondition/datamix_gemma/dataset_builders/orca_math_dataset_builder.py @@ -1,4 +1,4 @@ -# Copyright 2024 The precondition Authors. +# Copyright 2025 The precondition 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/precondition/datamix_gemma/dataset_builders/preprocessed_codealpaca_dataset_builder.py b/precondition/datamix_gemma/dataset_builders/preprocessed_codealpaca_dataset_builder.py index 42bebd0..d6ec872 100644 --- a/precondition/datamix_gemma/dataset_builders/preprocessed_codealpaca_dataset_builder.py +++ b/precondition/datamix_gemma/dataset_builders/preprocessed_codealpaca_dataset_builder.py @@ -1,4 +1,4 @@ -# Copyright 2024 The precondition Authors. +# Copyright 2025 The precondition 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/precondition/datamix_gemma/dataset_builders/preprocessed_dolly_dataset_builder.py b/precondition/datamix_gemma/dataset_builders/preprocessed_dolly_dataset_builder.py index d0d1ff2..0ecece1 100644 --- a/precondition/datamix_gemma/dataset_builders/preprocessed_dolly_dataset_builder.py +++ b/precondition/datamix_gemma/dataset_builders/preprocessed_dolly_dataset_builder.py @@ -1,4 +1,4 @@ -# Copyright 2024 The precondition Authors. +# Copyright 2025 The precondition 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/precondition/datamix_gemma/dataset_builders/preprocessed_gsm8k_dataset_builder.py b/precondition/datamix_gemma/dataset_builders/preprocessed_gsm8k_dataset_builder.py index dacf265..f676765 100644 --- a/precondition/datamix_gemma/dataset_builders/preprocessed_gsm8k_dataset_builder.py +++ b/precondition/datamix_gemma/dataset_builders/preprocessed_gsm8k_dataset_builder.py @@ -1,4 +1,4 @@ -# Copyright 2024 The precondition Authors. +# Copyright 2025 The precondition 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/precondition/datamix_gemma/dataset_builders/preprocessed_metamath_dataset_builder.py b/precondition/datamix_gemma/dataset_builders/preprocessed_metamath_dataset_builder.py index 3387070..9ae7899 100644 --- a/precondition/datamix_gemma/dataset_builders/preprocessed_metamath_dataset_builder.py +++ b/precondition/datamix_gemma/dataset_builders/preprocessed_metamath_dataset_builder.py @@ -1,4 +1,4 @@ -# Copyright 2024 The precondition Authors. +# Copyright 2025 The precondition 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/precondition/datamix_gemma/dataset_builders/preprocessed_open_orca_dataset_builder.py b/precondition/datamix_gemma/dataset_builders/preprocessed_open_orca_dataset_builder.py index 6842ce8..ccdb45c 100644 --- a/precondition/datamix_gemma/dataset_builders/preprocessed_open_orca_dataset_builder.py +++ b/precondition/datamix_gemma/dataset_builders/preprocessed_open_orca_dataset_builder.py @@ -1,4 +1,4 @@ -# Copyright 2024 The precondition Authors. +# Copyright 2025 The precondition 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/precondition/datamix_gemma/dataset_builders/preprocessed_orca_math_dataset_builder.py b/precondition/datamix_gemma/dataset_builders/preprocessed_orca_math_dataset_builder.py index 4dd87c3..51dc7f4 100644 --- a/precondition/datamix_gemma/dataset_builders/preprocessed_orca_math_dataset_builder.py +++ b/precondition/datamix_gemma/dataset_builders/preprocessed_orca_math_dataset_builder.py @@ -1,4 +1,4 @@ -# Copyright 2024 The precondition Authors. +# Copyright 2025 The precondition 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/precondition/datamix_gemma/dataset_builders/preprocessed_sciq_dataset_builder.py b/precondition/datamix_gemma/dataset_builders/preprocessed_sciq_dataset_builder.py index 0afceb0..9fdf510 100644 --- a/precondition/datamix_gemma/dataset_builders/preprocessed_sciq_dataset_builder.py +++ b/precondition/datamix_gemma/dataset_builders/preprocessed_sciq_dataset_builder.py @@ -1,4 +1,4 @@ -# Copyright 2024 The precondition Authors. +# Copyright 2025 The precondition 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/precondition/datamix_gemma/dataset_builders/preprocessed_wikipedia_dataset_builder.py b/precondition/datamix_gemma/dataset_builders/preprocessed_wikipedia_dataset_builder.py index 230331f..92a7156 100644 --- a/precondition/datamix_gemma/dataset_builders/preprocessed_wikipedia_dataset_builder.py +++ b/precondition/datamix_gemma/dataset_builders/preprocessed_wikipedia_dataset_builder.py @@ -1,4 +1,4 @@ -# Copyright 2024 The precondition Authors. +# Copyright 2025 The precondition 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/precondition/datamix_gemma/deconstructed_sampler.py b/precondition/datamix_gemma/deconstructed_sampler.py index a8dbb9d..e27ca10 100644 --- a/precondition/datamix_gemma/deconstructed_sampler.py +++ b/precondition/datamix_gemma/deconstructed_sampler.py @@ -1,4 +1,4 @@ -# Copyright 2024 The precondition Authors. +# Copyright 2025 The precondition 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/precondition/datamix_gemma/deterministic_strategy_bandit_loop.py b/precondition/datamix_gemma/deterministic_strategy_bandit_loop.py index fd41e80..e9a9ca6 100644 --- a/precondition/datamix_gemma/deterministic_strategy_bandit_loop.py +++ b/precondition/datamix_gemma/deterministic_strategy_bandit_loop.py @@ -1,4 +1,4 @@ -# Copyright 2024 The precondition Authors. +# Copyright 2025 The precondition 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/precondition/datamix_gemma/evals/__init__.py b/precondition/datamix_gemma/evals/__init__.py index 44d891e..7ee5dd2 100644 --- a/precondition/datamix_gemma/evals/__init__.py +++ b/precondition/datamix_gemma/evals/__init__.py @@ -1,4 +1,4 @@ -# Copyright 2024 The precondition Authors. +# Copyright 2025 The precondition 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/precondition/datamix_gemma/evals/crop.py b/precondition/datamix_gemma/evals/crop.py index 614a107..81c1966 100644 --- a/precondition/datamix_gemma/evals/crop.py +++ b/precondition/datamix_gemma/evals/crop.py @@ -1,4 +1,4 @@ -# Copyright 2024 The precondition Authors. +# Copyright 2025 The precondition 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/precondition/datamix_gemma/evals/eval.py b/precondition/datamix_gemma/evals/eval.py index af69f93..de3e0ec 100644 --- a/precondition/datamix_gemma/evals/eval.py +++ b/precondition/datamix_gemma/evals/eval.py @@ -1,4 +1,4 @@ -# Copyright 2024 The precondition Authors. +# Copyright 2025 The precondition 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/precondition/datamix_gemma/evals/gsm8k_eval.py b/precondition/datamix_gemma/evals/gsm8k_eval.py index 62834d0..c10d017 100644 --- a/precondition/datamix_gemma/evals/gsm8k_eval.py +++ b/precondition/datamix_gemma/evals/gsm8k_eval.py @@ -1,4 +1,4 @@ -# Copyright 2024 The precondition Authors. +# Copyright 2025 The precondition 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/precondition/datamix_gemma/evals/mbpp_eval.py b/precondition/datamix_gemma/evals/mbpp_eval.py index 920e49c..4707150 100644 --- a/precondition/datamix_gemma/evals/mbpp_eval.py +++ b/precondition/datamix_gemma/evals/mbpp_eval.py @@ -1,4 +1,4 @@ -# Copyright 2024 The precondition Authors. +# Copyright 2025 The precondition 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/precondition/datamix_gemma/evals/mmlu_eval.py b/precondition/datamix_gemma/evals/mmlu_eval.py index ddc7720..7bbae45 100644 --- a/precondition/datamix_gemma/evals/mmlu_eval.py +++ b/precondition/datamix_gemma/evals/mmlu_eval.py @@ -1,4 +1,4 @@ -# Copyright 2024 The precondition Authors. +# Copyright 2025 The precondition 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/precondition/datamix_gemma/finetune.py b/precondition/datamix_gemma/finetune.py index ebf66f7..e87c009 100644 --- a/precondition/datamix_gemma/finetune.py +++ b/precondition/datamix_gemma/finetune.py @@ -1,4 +1,4 @@ -# Copyright 2024 The precondition Authors. +# Copyright 2025 The precondition 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/precondition/datamix_gemma/finetune_eval_measurement.py b/precondition/datamix_gemma/finetune_eval_measurement.py index c0cf982..77e5ce4 100644 --- a/precondition/datamix_gemma/finetune_eval_measurement.py +++ b/precondition/datamix_gemma/finetune_eval_measurement.py @@ -1,4 +1,4 @@ -# Copyright 2024 The precondition Authors. +# Copyright 2025 The precondition 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/precondition/datamix_gemma/finetune_utils.py b/precondition/datamix_gemma/finetune_utils.py index bb8d33d..fc4c64b 100644 --- a/precondition/datamix_gemma/finetune_utils.py +++ b/precondition/datamix_gemma/finetune_utils.py @@ -1,4 +1,4 @@ -# Copyright 2024 The precondition Authors. +# Copyright 2025 The precondition 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/precondition/datamix_gemma/finetuning_experiment.py b/precondition/datamix_gemma/finetuning_experiment.py index b634930..ba1a707 100644 --- a/precondition/datamix_gemma/finetuning_experiment.py +++ b/precondition/datamix_gemma/finetuning_experiment.py @@ -1,4 +1,4 @@ -# Copyright 2024 The precondition Authors. +# Copyright 2025 The precondition 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/precondition/datamix_gemma/random_baseline.py b/precondition/datamix_gemma/random_baseline.py index 66cf4d3..c8040a6 100644 --- a/precondition/datamix_gemma/random_baseline.py +++ b/precondition/datamix_gemma/random_baseline.py @@ -1,4 +1,4 @@ -# Copyright 2024 The precondition Authors. +# Copyright 2025 The precondition 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/precondition/datamix_gemma/snr_calculation.py b/precondition/datamix_gemma/snr_calculation.py index bc63ac7..98151d9 100644 --- a/precondition/datamix_gemma/snr_calculation.py +++ b/precondition/datamix_gemma/snr_calculation.py @@ -1,4 +1,4 @@ -# Copyright 2024 The precondition Authors. +# Copyright 2025 The precondition 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/precondition/datamix_gemma/tokenizers/__init__.py b/precondition/datamix_gemma/tokenizers/__init__.py index 44d891e..7ee5dd2 100644 --- a/precondition/datamix_gemma/tokenizers/__init__.py +++ b/precondition/datamix_gemma/tokenizers/__init__.py @@ -1,4 +1,4 @@ -# Copyright 2024 The precondition Authors. +# Copyright 2025 The precondition 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/precondition/datamix_gemma/tokenizers/gemma_tokenizer.py b/precondition/datamix_gemma/tokenizers/gemma_tokenizer.py index 41219b4..3cdd69d 100644 --- a/precondition/datamix_gemma/tokenizers/gemma_tokenizer.py +++ b/precondition/datamix_gemma/tokenizers/gemma_tokenizer.py @@ -1,4 +1,4 @@ -# Copyright 2024 The precondition Authors. +# Copyright 2025 The precondition 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/precondition/datamix_gemma/training_batch_generators/__init__.py b/precondition/datamix_gemma/training_batch_generators/__init__.py index 44d891e..7ee5dd2 100644 --- a/precondition/datamix_gemma/training_batch_generators/__init__.py +++ b/precondition/datamix_gemma/training_batch_generators/__init__.py @@ -1,4 +1,4 @@ -# Copyright 2024 The precondition Authors. +# Copyright 2025 The precondition 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/precondition/datamix_gemma/training_batch_generators/dartboard_deterministic_training_batch_generator.py b/precondition/datamix_gemma/training_batch_generators/dartboard_deterministic_training_batch_generator.py index ac3f462..adc4cff 100644 --- a/precondition/datamix_gemma/training_batch_generators/dartboard_deterministic_training_batch_generator.py +++ b/precondition/datamix_gemma/training_batch_generators/dartboard_deterministic_training_batch_generator.py @@ -1,4 +1,4 @@ -# Copyright 2024 The precondition Authors. +# Copyright 2025 The precondition 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/precondition/datamix_gemma/training_batch_generators/dartboard_importance_weighting_training_batch_generator.py b/precondition/datamix_gemma/training_batch_generators/dartboard_importance_weighting_training_batch_generator.py index 1d49249..505f5a7 100644 --- a/precondition/datamix_gemma/training_batch_generators/dartboard_importance_weighting_training_batch_generator.py +++ b/precondition/datamix_gemma/training_batch_generators/dartboard_importance_weighting_training_batch_generator.py @@ -1,4 +1,4 @@ -# Copyright 2024 The precondition Authors. +# Copyright 2025 The precondition 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/precondition/datamix_gemma/training_batch_generators/fixed_dataset_importance_weighting_training_batch_generator.py b/precondition/datamix_gemma/training_batch_generators/fixed_dataset_importance_weighting_training_batch_generator.py index c7b2b3d..a6b75a6 100644 --- a/precondition/datamix_gemma/training_batch_generators/fixed_dataset_importance_weighting_training_batch_generator.py +++ b/precondition/datamix_gemma/training_batch_generators/fixed_dataset_importance_weighting_training_batch_generator.py @@ -1,4 +1,4 @@ -# Copyright 2024 The precondition Authors. +# Copyright 2025 The precondition 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/precondition/datamix_gemma/training_batch_generators/importance_weighting_training_batch_generator.py b/precondition/datamix_gemma/training_batch_generators/importance_weighting_training_batch_generator.py index b418434..5d078c7 100644 --- a/precondition/datamix_gemma/training_batch_generators/importance_weighting_training_batch_generator.py +++ b/precondition/datamix_gemma/training_batch_generators/importance_weighting_training_batch_generator.py @@ -1,4 +1,4 @@ -# Copyright 2024 The precondition Authors. +# Copyright 2025 The precondition 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/precondition/datamix_gemma/training_batch_generators/training_batch_generator.py b/precondition/datamix_gemma/training_batch_generators/training_batch_generator.py index 68d6164..bc64170 100644 --- a/precondition/datamix_gemma/training_batch_generators/training_batch_generator.py +++ b/precondition/datamix_gemma/training_batch_generators/training_batch_generator.py @@ -1,4 +1,4 @@ -# Copyright 2024 The precondition Authors. +# Copyright 2025 The precondition 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/precondition/datamix_gemma/training_batch_generators/vanilla_training_batch_generator.py b/precondition/datamix_gemma/training_batch_generators/vanilla_training_batch_generator.py index 4bf8a0d..ef34137 100644 --- a/precondition/datamix_gemma/training_batch_generators/vanilla_training_batch_generator.py +++ b/precondition/datamix_gemma/training_batch_generators/vanilla_training_batch_generator.py @@ -1,4 +1,4 @@ -# Copyright 2024 The precondition Authors. +# Copyright 2025 The precondition 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/precondition/datamix_gemma/training_loop.py b/precondition/datamix_gemma/training_loop.py index db30e19..b73958e 100644 --- a/precondition/datamix_gemma/training_loop.py +++ b/precondition/datamix_gemma/training_loop.py @@ -1,4 +1,4 @@ -# Copyright 2024 The precondition Authors. +# Copyright 2025 The precondition 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/precondition/distributed_shampoo.py b/precondition/distributed_shampoo.py index 8ad9655..c3fce10 100644 --- a/precondition/distributed_shampoo.py +++ b/precondition/distributed_shampoo.py @@ -1,4 +1,4 @@ -# Copyright 2024 The precondition Authors. +# Copyright 2025 The precondition Authors. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -649,7 +649,7 @@ def _iter_body(state): _, v_out, s_out, _, _ = lax.while_loop(_iter_condition, _iter_body, init_state) v_out = v_out / jnp.linalg.norm(v_out) - return v_out, s_out + return v_out, s_out # pytype: disable=bad-return-type # lax-types def mat_power( @@ -876,7 +876,7 @@ def _outer_body_fn(state): iters, mat_m, mat_h, old_mat_h, error, error_ratio = lax.while_loop( _iter_condition, _iter_body, init_state) error = jnp.max(jnp.abs(mat_m - identity)).astype(jnp.float32) - is_converged = jnp.asarray(error_ratio < max_error_ratio, old_mat_h.dtype) + is_converged = jnp.asarray(error_ratio < max_error_ratio, old_mat_h.dtype) # pytype: disable=attribute-error # lax-types resultant_mat_h = is_converged * mat_h + (1 - is_converged) * old_mat_h return (i + 1, resultant_mat_h, error, iters, error_ratio, error > retry_loop_error_threshold) @@ -2871,7 +2871,7 @@ def _internal_inverse_pth_root_all(): all_statistics[current_replica], all_exponents[current_replica], all_paddings[current_replica], - _maybe_ix(all_preconditioners, current_replica), + _maybe_ix(all_preconditioners, current_replica), # pytype: disable=wrong-arg-types # lax-types ) preconditioners = jax.lax.all_gather(preconditioners, batch_axis_name) metrics = jax.lax.all_gather(metrics, batch_axis_name) @@ -3109,9 +3109,9 @@ def _internal_inverse_pth_root_all(): all_quantized_bucket_sizes[current_replica], all_exponents[current_replica], all_paddings[current_replica], - _maybe_ix(all_quantized_precond_mats, current_replica), - _maybe_ix(all_quantized_precond_diagonals, current_replica), - _maybe_ix(all_quantized_precond_bucket_sizes, current_replica), + _maybe_ix(all_quantized_precond_mats, current_replica), # pytype: disable=wrong-arg-types # lax-types + _maybe_ix(all_quantized_precond_diagonals, current_replica), # pytype: disable=wrong-arg-types # lax-types + _maybe_ix(all_quantized_precond_bucket_sizes, current_replica), # pytype: disable=wrong-arg-types # lax-types ) quantized_preconditioners = jax.lax.all_gather(quantized_preconditioners, batch_axis_name) diff --git a/precondition/distributed_shampoo_test.py b/precondition/distributed_shampoo_test.py index 4196348..0d93bcb 100644 --- a/precondition/distributed_shampoo_test.py +++ b/precondition/distributed_shampoo_test.py @@ -1,4 +1,4 @@ -# Copyright 2024 The precondition Authors. +# Copyright 2025 The precondition 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/precondition/oco/algorithms.py b/precondition/oco/algorithms.py index 16a5b53..7bdbf42 100644 --- a/precondition/oco/algorithms.py +++ b/precondition/oco/algorithms.py @@ -1,4 +1,4 @@ -# Copyright 2024 The precondition Authors. +# Copyright 2025 The precondition 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/precondition/oco/datasets.py b/precondition/oco/datasets.py index 20e2dfe..639abf2 100644 --- a/precondition/oco/datasets.py +++ b/precondition/oco/datasets.py @@ -1,4 +1,4 @@ -# Copyright 2024 The precondition Authors. +# Copyright 2025 The precondition 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/precondition/oco/sweep.py b/precondition/oco/sweep.py index 0baa9b5..6b8f7bd 100644 --- a/precondition/oco/sweep.py +++ b/precondition/oco/sweep.py @@ -1,4 +1,4 @@ -# Copyright 2024 The precondition Authors. +# Copyright 2025 The precondition 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/precondition/oco/train.py b/precondition/oco/train.py index 34cccf8..33a4e3b 100644 --- a/precondition/oco/train.py +++ b/precondition/oco/train.py @@ -1,4 +1,4 @@ -# Copyright 2024 The precondition Authors. +# Copyright 2025 The precondition 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/precondition/quantization_utils.py b/precondition/quantization_utils.py index d721df3..58ea590 100644 --- a/precondition/quantization_utils.py +++ b/precondition/quantization_utils.py @@ -1,4 +1,4 @@ -# Copyright 2024 The precondition Authors. +# Copyright 2025 The precondition 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/precondition/sm3.py b/precondition/sm3.py index d44ff74..ad90c28 100644 --- a/precondition/sm3.py +++ b/precondition/sm3.py @@ -1,4 +1,4 @@ -# Copyright 2024 The precondition Authors. +# Copyright 2025 The precondition 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/precondition/sm3_test.py b/precondition/sm3_test.py index 95814cd..d6bf658 100644 --- a/precondition/sm3_test.py +++ b/precondition/sm3_test.py @@ -1,4 +1,4 @@ -# Copyright 2024 The precondition Authors. +# Copyright 2025 The precondition 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/precondition/tearfree/grafting.py b/precondition/tearfree/grafting.py index 18acb54..72c9305 100644 --- a/precondition/tearfree/grafting.py +++ b/precondition/tearfree/grafting.py @@ -1,4 +1,4 @@ -# Copyright 2024 The precondition Authors. +# Copyright 2025 The precondition 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/precondition/tearfree/grafting_test.py b/precondition/tearfree/grafting_test.py index 0d866f0..56ec36f 100644 --- a/precondition/tearfree/grafting_test.py +++ b/precondition/tearfree/grafting_test.py @@ -1,4 +1,4 @@ -# Copyright 2024 The precondition Authors. +# Copyright 2025 The precondition 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/precondition/tearfree/momentum.py b/precondition/tearfree/momentum.py index 67b58c6..c01f232 100644 --- a/precondition/tearfree/momentum.py +++ b/precondition/tearfree/momentum.py @@ -1,4 +1,4 @@ -# Copyright 2024 The precondition Authors. +# Copyright 2025 The precondition 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/precondition/tearfree/momentum_test.py b/precondition/tearfree/momentum_test.py index 2dab345..a4715de 100644 --- a/precondition/tearfree/momentum_test.py +++ b/precondition/tearfree/momentum_test.py @@ -1,4 +1,4 @@ -# Copyright 2024 The precondition Authors. +# Copyright 2025 The precondition 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/precondition/tearfree/optimizer.py b/precondition/tearfree/optimizer.py index 759059f..afbf814 100644 --- a/precondition/tearfree/optimizer.py +++ b/precondition/tearfree/optimizer.py @@ -1,4 +1,4 @@ -# Copyright 2024 The precondition Authors. +# Copyright 2025 The precondition 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/precondition/tearfree/optimizer_smoke_test.py b/precondition/tearfree/optimizer_smoke_test.py index a259986..38400c7 100644 --- a/precondition/tearfree/optimizer_smoke_test.py +++ b/precondition/tearfree/optimizer_smoke_test.py @@ -1,4 +1,4 @@ -# Copyright 2024 The precondition Authors. +# Copyright 2025 The precondition 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/precondition/tearfree/optimizer_test.py b/precondition/tearfree/optimizer_test.py index 2d3ed2d..3c5164f 100644 --- a/precondition/tearfree/optimizer_test.py +++ b/precondition/tearfree/optimizer_test.py @@ -1,4 +1,4 @@ -# Copyright 2024 The precondition Authors. +# Copyright 2025 The precondition 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/precondition/tearfree/praxis_shim.py b/precondition/tearfree/praxis_shim.py index 96dd36b..d6c4c7d 100644 --- a/precondition/tearfree/praxis_shim.py +++ b/precondition/tearfree/praxis_shim.py @@ -1,4 +1,4 @@ -# Copyright 2024 The precondition Authors. +# Copyright 2025 The precondition 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/precondition/tearfree/reallocation.py b/precondition/tearfree/reallocation.py index 355faed..d4688f1 100644 --- a/precondition/tearfree/reallocation.py +++ b/precondition/tearfree/reallocation.py @@ -1,4 +1,4 @@ -# Copyright 2024 The precondition Authors. +# Copyright 2025 The precondition 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/precondition/tearfree/reallocation_test.py b/precondition/tearfree/reallocation_test.py index 7d87f2a..701784e 100644 --- a/precondition/tearfree/reallocation_test.py +++ b/precondition/tearfree/reallocation_test.py @@ -1,4 +1,4 @@ -# Copyright 2024 The precondition Authors. +# Copyright 2025 The precondition 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/precondition/tearfree/reshaper.py b/precondition/tearfree/reshaper.py index 3e9efa8..a983574 100644 --- a/precondition/tearfree/reshaper.py +++ b/precondition/tearfree/reshaper.py @@ -1,4 +1,4 @@ -# Copyright 2024 The precondition Authors. +# Copyright 2025 The precondition 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/precondition/tearfree/reshaper_test.py b/precondition/tearfree/reshaper_test.py index f3c41a4..cb90e02 100644 --- a/precondition/tearfree/reshaper_test.py +++ b/precondition/tearfree/reshaper_test.py @@ -1,4 +1,4 @@ -# Copyright 2024 The precondition Authors. +# Copyright 2025 The precondition 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/precondition/tearfree/second_order.py b/precondition/tearfree/second_order.py index 63f1e2f..62cc6de 100644 --- a/precondition/tearfree/second_order.py +++ b/precondition/tearfree/second_order.py @@ -1,4 +1,4 @@ -# Copyright 2024 The precondition Authors. +# Copyright 2025 The precondition 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/precondition/tearfree/shampoo.py b/precondition/tearfree/shampoo.py index 86526b1..434f913 100644 --- a/precondition/tearfree/shampoo.py +++ b/precondition/tearfree/shampoo.py @@ -1,4 +1,4 @@ -# Copyright 2024 The precondition Authors. +# Copyright 2025 The precondition 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/precondition/tearfree/shampoo_test.py b/precondition/tearfree/shampoo_test.py index db81222..e6705a5 100644 --- a/precondition/tearfree/shampoo_test.py +++ b/precondition/tearfree/shampoo_test.py @@ -1,4 +1,4 @@ -# Copyright 2024 The precondition Authors. +# Copyright 2025 The precondition 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/precondition/tearfree/sketchy.py b/precondition/tearfree/sketchy.py index 37bf1bb..9e7e1de 100644 --- a/precondition/tearfree/sketchy.py +++ b/precondition/tearfree/sketchy.py @@ -1,4 +1,4 @@ -# Copyright 2024 The precondition Authors. +# Copyright 2025 The precondition 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/precondition/tearfree/sketchy_test.py b/precondition/tearfree/sketchy_test.py index a39d16b..1762e7a 100644 --- a/precondition/tearfree/sketchy_test.py +++ b/precondition/tearfree/sketchy_test.py @@ -1,4 +1,4 @@ -# Copyright 2024 The precondition Authors. +# Copyright 2025 The precondition Authors. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License.