From 0c108255bf3d5e59eed578cb1a5a7ce22b8fd0fc Mon Sep 17 00:00:00 2001 From: Joan Puigcerver Date: Fri, 20 Jun 2025 00:36:19 -0700 Subject: [PATCH] Minor fix after jax changes. PiperOrigin-RevId: 773571781 --- vmoe/app.py | 2 +- vmoe/checkpoints/__init__.py | 2 +- vmoe/checkpoints/base.py | 2 +- vmoe/checkpoints/base_test.py | 2 +- vmoe/checkpoints/partitioned.py | 8 +++++--- vmoe/checkpoints/partitioned_test.py | 2 +- vmoe/checkpoints/serialization.py | 2 +- vmoe/checkpoints/serialization_test.py | 2 +- vmoe/checkpoints/types.py | 2 +- vmoe/checkpoints/types_test.py | 2 +- vmoe/configs/common_fewshot.py | 2 +- .../eee_paper/eee_s32_last2_ilsvrc2012_ft_cifar100.py | 2 +- vmoe/configs/vmoe_paper/common.py | 2 +- vmoe/configs/vmoe_paper/pretrain_imagenet21k.py | 2 +- .../vmoe_b16_imagenet21k_randaug_strong_ft_cifar10.py | 2 +- .../vmoe_b16_imagenet21k_randaug_strong_ft_ilsvrc2012.py | 2 +- .../vmoe_s32_last2_ilsvrc2012_randaug_light1.py | 2 +- ...e_s32_last2_ilsvrc2012_randaug_light1_ft_ilsvrc2012.py | 2 +- vmoe/constants.py | 2 +- vmoe/data/builder.py | 2 +- vmoe/data/builder_test.py | 2 +- vmoe/data/input_pipeline.py | 2 +- vmoe/data/input_pipeline_test.py | 2 +- vmoe/data/pjit_utils.py | 2 +- vmoe/data/pjit_utils_test.py | 2 +- vmoe/data/pp_ops.py | 2 +- vmoe/data/pp_ops_test.py | 2 +- vmoe/evaluate/ensemble.py | 2 +- vmoe/evaluate/ensemble_test.py | 2 +- vmoe/evaluate/evaluator.py | 2 +- vmoe/evaluate/evaluator_test.py | 2 +- vmoe/evaluate/fewshot.py | 2 +- vmoe/evaluate/fewshot_test.py | 2 +- vmoe/initialization/__init__.py | 2 +- vmoe/initialization/initialization.py | 2 +- vmoe/initialization/initialization_test.py | 2 +- vmoe/initialization/mapping.py | 2 +- vmoe/initialization/mapping_test.py | 2 +- vmoe/initialization/rules.py | 2 +- vmoe/initialization/rules_test.py | 2 +- vmoe/moe.py | 2 +- vmoe/moe_test.py | 2 +- vmoe/multihost_utils.py | 2 +- vmoe/nn/ensemble_routing.py | 2 +- vmoe/nn/ensemble_routing_test.py | 2 +- vmoe/nn/external.py | 2 +- vmoe/nn/external_test.py | 2 +- vmoe/nn/models.py | 2 +- vmoe/nn/routing.py | 2 +- vmoe/nn/routing_test.py | 2 +- vmoe/nn/vit_moe.py | 2 +- vmoe/nn/vit_moe_ensemble.py | 2 +- vmoe/nn/vit_moe_ensemble_test.py | 2 +- vmoe/nn/vit_moe_test.py | 2 +- vmoe/partitioning.py | 4 ++-- vmoe/partitioning_test.py | 2 +- vmoe/projects/adversarial_attacks/attacks.py | 2 +- vmoe/projects/adversarial_attacks/attacks_test.py | 4 ++-- .../adversarial_attacks/configs/attack/ilsvrc2012.py | 2 +- vmoe/projects/adversarial_attacks/configs/common.py | 2 +- vmoe/projects/adversarial_attacks/lib.py | 2 +- vmoe/projects/adversarial_attacks/lib_test.py | 2 +- vmoe/projects/adversarial_attacks/main.py | 2 +- vmoe/projects/adversarial_attacks/restore.py | 2 +- vmoe/projects/adversarial_attacks/restore_test.py | 2 +- vmoe/projects/contrastive/evaluators.py | 2 +- vmoe/projects/contrastive/models.py | 2 +- vmoe/projects/contrastive/models_test.py | 2 +- vmoe/projects/contrastive/trainer.py | 2 +- vmoe/projects/soft_moe/configs/common.py | 2 +- vmoe/projects/soft_moe/configs/pretrain_jft4b.py | 2 +- vmoe/projects/soft_moe/configs/pretrain_laion.py | 2 +- vmoe/projects/soft_moe/main.py | 2 +- vmoe/projects/soft_moe/main_contrastive.py | 2 +- vmoe/projects/soft_moe/router.py | 2 +- vmoe/projects/soft_moe/router_test.py | 2 +- .../sparsity_constrained_ot/kl_projection_routing.py | 2 +- .../sparsity_constrained_ot/kl_projection_routing_test.py | 2 +- .../sparsity_constrained_ot/ksparse_projection_routing.py | 2 +- .../ksparse_projection_routing_test.py | 2 +- vmoe/projects/sparsity_constrained_ot/ot_routing.py | 2 +- .../sparsity_constrained_ot/sparse_projection_routing.py | 2 +- .../sparse_projection_routing_test.py | 2 +- vmoe/train/main.py | 2 +- vmoe/train/optimizer.py | 2 +- vmoe/train/optimizer_test.py | 2 +- vmoe/train/periodic_actions.py | 2 +- vmoe/train/periodic_actions_test.py | 2 +- vmoe/train/schedule.py | 2 +- vmoe/train/schedule_test.py | 2 +- vmoe/train/train_state.py | 2 +- vmoe/train/train_state_test.py | 2 +- vmoe/train/trainer.py | 2 +- vmoe/train/trainer_test.py | 2 +- vmoe/train/tree_summarizer.py | 2 +- vmoe/train/tree_summarizer_test.py | 2 +- vmoe/utils.py | 2 +- vmoe/utils_test.py | 2 +- 98 files changed, 104 insertions(+), 102 deletions(-) diff --git a/vmoe/app.py b/vmoe/app.py index 7915b77..5b0c637 100644 --- a/vmoe/app.py +++ b/vmoe/app.py @@ -1,4 +1,4 @@ -# Copyright 2024 Google LLC. +# Copyright 2025 Google LLC. # # 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/vmoe/checkpoints/__init__.py b/vmoe/checkpoints/__init__.py index 1e3ff7f..637e06f 100644 --- a/vmoe/checkpoints/__init__.py +++ b/vmoe/checkpoints/__init__.py @@ -1,4 +1,4 @@ -# Copyright 2024 Google LLC. +# Copyright 2025 Google LLC. # # 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/vmoe/checkpoints/base.py b/vmoe/checkpoints/base.py index 096f811..4a24c2a 100644 --- a/vmoe/checkpoints/base.py +++ b/vmoe/checkpoints/base.py @@ -1,4 +1,4 @@ -# Copyright 2024 Google LLC. +# Copyright 2025 Google LLC. # # 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/vmoe/checkpoints/base_test.py b/vmoe/checkpoints/base_test.py index 8e6bf40..fad446c 100644 --- a/vmoe/checkpoints/base_test.py +++ b/vmoe/checkpoints/base_test.py @@ -1,4 +1,4 @@ -# Copyright 2024 Google LLC. +# Copyright 2025 Google LLC. # # 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/vmoe/checkpoints/partitioned.py b/vmoe/checkpoints/partitioned.py index 99cf06d..8fe799e 100644 --- a/vmoe/checkpoints/partitioned.py +++ b/vmoe/checkpoints/partitioned.py @@ -1,4 +1,4 @@ -# Copyright 2024 Google LLC. +# Copyright 2025 Google LLC. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -56,8 +56,10 @@ IndexInfo = types.IndexInfo LazyArrayChunks = types.LazyArrayChunks MapResult = multiprocessing.pool.MapResult -GSPMDSharding = jax.sharding.GSPMDSharding PyTree = Any +Mesh = jax.sharding.Mesh +NamedSharding = jax.sharding.NamedSharding +PartitionSpec = jax.sharding.PartitionSpec Sharding = jax.sharding.Sharding Slice = types.Slice SliceNd = types.SliceNd @@ -361,7 +363,7 @@ def _get_array_sharding_or_default(arr: jax.Array) -> Sharding: if hasattr(arr, 'sharding'): return arr.sharding else: - return GSPMDSharding.get_replicated(jax.devices()) + return NamedSharding(Mesh(jax.devices(), 'x'), PartitionSpec()) def _intersect_slicend(a: SliceNd, b: SliceNd) -> Optional[SliceNd]: diff --git a/vmoe/checkpoints/partitioned_test.py b/vmoe/checkpoints/partitioned_test.py index ab1636c..920b87b 100644 --- a/vmoe/checkpoints/partitioned_test.py +++ b/vmoe/checkpoints/partitioned_test.py @@ -1,4 +1,4 @@ -# Copyright 2024 Google LLC. +# Copyright 2025 Google LLC. # # 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/vmoe/checkpoints/serialization.py b/vmoe/checkpoints/serialization.py index 3144e40..438e1db 100644 --- a/vmoe/checkpoints/serialization.py +++ b/vmoe/checkpoints/serialization.py @@ -1,4 +1,4 @@ -# Copyright 2024 Google LLC. +# Copyright 2025 Google LLC. # # 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/vmoe/checkpoints/serialization_test.py b/vmoe/checkpoints/serialization_test.py index 0d731a6..ab789cd 100644 --- a/vmoe/checkpoints/serialization_test.py +++ b/vmoe/checkpoints/serialization_test.py @@ -1,4 +1,4 @@ -# Copyright 2024 Google LLC. +# Copyright 2025 Google LLC. # # 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/vmoe/checkpoints/types.py b/vmoe/checkpoints/types.py index f0ad5c3..98dadb7 100644 --- a/vmoe/checkpoints/types.py +++ b/vmoe/checkpoints/types.py @@ -1,4 +1,4 @@ -# Copyright 2024 Google LLC. +# Copyright 2025 Google LLC. # # 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/vmoe/checkpoints/types_test.py b/vmoe/checkpoints/types_test.py index 2342937..8102090 100644 --- a/vmoe/checkpoints/types_test.py +++ b/vmoe/checkpoints/types_test.py @@ -1,4 +1,4 @@ -# Copyright 2024 Google LLC. +# Copyright 2025 Google LLC. # # 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/vmoe/configs/common_fewshot.py b/vmoe/configs/common_fewshot.py index 9bd1b87..c2cd475 100644 --- a/vmoe/configs/common_fewshot.py +++ b/vmoe/configs/common_fewshot.py @@ -1,4 +1,4 @@ -# Copyright 2024 Google LLC. +# Copyright 2025 Google LLC. # # 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/vmoe/configs/eee_paper/eee_s32_last2_ilsvrc2012_ft_cifar100.py b/vmoe/configs/eee_paper/eee_s32_last2_ilsvrc2012_ft_cifar100.py index 7cad8f3..0170ecf 100644 --- a/vmoe/configs/eee_paper/eee_s32_last2_ilsvrc2012_ft_cifar100.py +++ b/vmoe/configs/eee_paper/eee_s32_last2_ilsvrc2012_ft_cifar100.py @@ -1,4 +1,4 @@ -# Copyright 2024 Google LLC. +# Copyright 2025 Google LLC. # # 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/vmoe/configs/vmoe_paper/common.py b/vmoe/configs/vmoe_paper/common.py index 8e1644c..f9e42a0 100644 --- a/vmoe/configs/vmoe_paper/common.py +++ b/vmoe/configs/vmoe_paper/common.py @@ -1,4 +1,4 @@ -# Copyright 2024 Google LLC. +# Copyright 2025 Google LLC. # # 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/vmoe/configs/vmoe_paper/pretrain_imagenet21k.py b/vmoe/configs/vmoe_paper/pretrain_imagenet21k.py index e3230e1..be9cc30 100644 --- a/vmoe/configs/vmoe_paper/pretrain_imagenet21k.py +++ b/vmoe/configs/vmoe_paper/pretrain_imagenet21k.py @@ -1,4 +1,4 @@ -# Copyright 2024 Google LLC. +# Copyright 2025 Google LLC. # # 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/vmoe/configs/vmoe_paper/vmoe_b16_imagenet21k_randaug_strong_ft_cifar10.py b/vmoe/configs/vmoe_paper/vmoe_b16_imagenet21k_randaug_strong_ft_cifar10.py index da428e0..635e181 100644 --- a/vmoe/configs/vmoe_paper/vmoe_b16_imagenet21k_randaug_strong_ft_cifar10.py +++ b/vmoe/configs/vmoe_paper/vmoe_b16_imagenet21k_randaug_strong_ft_cifar10.py @@ -1,4 +1,4 @@ -# Copyright 2024 Google LLC. +# Copyright 2025 Google LLC. # # 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/vmoe/configs/vmoe_paper/vmoe_b16_imagenet21k_randaug_strong_ft_ilsvrc2012.py b/vmoe/configs/vmoe_paper/vmoe_b16_imagenet21k_randaug_strong_ft_ilsvrc2012.py index e9ce7f4..96ef289 100644 --- a/vmoe/configs/vmoe_paper/vmoe_b16_imagenet21k_randaug_strong_ft_ilsvrc2012.py +++ b/vmoe/configs/vmoe_paper/vmoe_b16_imagenet21k_randaug_strong_ft_ilsvrc2012.py @@ -1,4 +1,4 @@ -# Copyright 2024 Google LLC. +# Copyright 2025 Google LLC. # # 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/vmoe/configs/vmoe_paper/vmoe_s32_last2_ilsvrc2012_randaug_light1.py b/vmoe/configs/vmoe_paper/vmoe_s32_last2_ilsvrc2012_randaug_light1.py index 4b922d8..6b791ab 100644 --- a/vmoe/configs/vmoe_paper/vmoe_s32_last2_ilsvrc2012_randaug_light1.py +++ b/vmoe/configs/vmoe_paper/vmoe_s32_last2_ilsvrc2012_randaug_light1.py @@ -1,4 +1,4 @@ -# Copyright 2024 Google LLC. +# Copyright 2025 Google LLC. # # 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/vmoe/configs/vmoe_paper/vmoe_s32_last2_ilsvrc2012_randaug_light1_ft_ilsvrc2012.py b/vmoe/configs/vmoe_paper/vmoe_s32_last2_ilsvrc2012_randaug_light1_ft_ilsvrc2012.py index 56109e4..07be2ab 100644 --- a/vmoe/configs/vmoe_paper/vmoe_s32_last2_ilsvrc2012_randaug_light1_ft_ilsvrc2012.py +++ b/vmoe/configs/vmoe_paper/vmoe_s32_last2_ilsvrc2012_randaug_light1_ft_ilsvrc2012.py @@ -1,4 +1,4 @@ -# Copyright 2024 Google LLC. +# Copyright 2025 Google LLC. # # 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/vmoe/constants.py b/vmoe/constants.py index 88d6bf8..d000a8a 100644 --- a/vmoe/constants.py +++ b/vmoe/constants.py @@ -1,4 +1,4 @@ -# Copyright 2024 Google LLC. +# Copyright 2025 Google LLC. # # 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/vmoe/data/builder.py b/vmoe/data/builder.py index 5e72f82..54e64b7 100644 --- a/vmoe/data/builder.py +++ b/vmoe/data/builder.py @@ -1,4 +1,4 @@ -# Copyright 2024 Google LLC. +# Copyright 2025 Google LLC. # # 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/vmoe/data/builder_test.py b/vmoe/data/builder_test.py index 9f10d49..c52e699 100644 --- a/vmoe/data/builder_test.py +++ b/vmoe/data/builder_test.py @@ -1,4 +1,4 @@ -# Copyright 2024 Google LLC. +# Copyright 2025 Google LLC. # # 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/vmoe/data/input_pipeline.py b/vmoe/data/input_pipeline.py index 8fb6b7b..613b73e 100644 --- a/vmoe/data/input_pipeline.py +++ b/vmoe/data/input_pipeline.py @@ -1,4 +1,4 @@ -# Copyright 2024 Google LLC. +# Copyright 2025 Google LLC. # # 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/vmoe/data/input_pipeline_test.py b/vmoe/data/input_pipeline_test.py index 3877a41..4714db1 100644 --- a/vmoe/data/input_pipeline_test.py +++ b/vmoe/data/input_pipeline_test.py @@ -1,4 +1,4 @@ -# Copyright 2024 Google LLC. +# Copyright 2025 Google LLC. # # 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/vmoe/data/pjit_utils.py b/vmoe/data/pjit_utils.py index 0640feb..0795a0a 100644 --- a/vmoe/data/pjit_utils.py +++ b/vmoe/data/pjit_utils.py @@ -1,4 +1,4 @@ -# Copyright 2024 Google LLC. +# Copyright 2025 Google LLC. # # 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/vmoe/data/pjit_utils_test.py b/vmoe/data/pjit_utils_test.py index bb3aa2a..7c21dc3 100644 --- a/vmoe/data/pjit_utils_test.py +++ b/vmoe/data/pjit_utils_test.py @@ -1,4 +1,4 @@ -# Copyright 2024 Google LLC. +# Copyright 2025 Google LLC. # # 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/vmoe/data/pp_ops.py b/vmoe/data/pp_ops.py index cb54efd..adde007 100644 --- a/vmoe/data/pp_ops.py +++ b/vmoe/data/pp_ops.py @@ -1,4 +1,4 @@ -# Copyright 2024 Google LLC. +# Copyright 2025 Google LLC. # # 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/vmoe/data/pp_ops_test.py b/vmoe/data/pp_ops_test.py index c35d01f..f9ec73e 100644 --- a/vmoe/data/pp_ops_test.py +++ b/vmoe/data/pp_ops_test.py @@ -1,4 +1,4 @@ -# Copyright 2024 Google LLC. +# Copyright 2025 Google LLC. # # 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/vmoe/evaluate/ensemble.py b/vmoe/evaluate/ensemble.py index b2678b9..bc6434e 100644 --- a/vmoe/evaluate/ensemble.py +++ b/vmoe/evaluate/ensemble.py @@ -1,4 +1,4 @@ -# Copyright 2024 Google LLC. +# Copyright 2025 Google LLC. # # 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/vmoe/evaluate/ensemble_test.py b/vmoe/evaluate/ensemble_test.py index d8b14a9..df2646c 100644 --- a/vmoe/evaluate/ensemble_test.py +++ b/vmoe/evaluate/ensemble_test.py @@ -1,4 +1,4 @@ -# Copyright 2024 Google LLC. +# Copyright 2025 Google LLC. # # 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/vmoe/evaluate/evaluator.py b/vmoe/evaluate/evaluator.py index 821b02f..176d273 100644 --- a/vmoe/evaluate/evaluator.py +++ b/vmoe/evaluate/evaluator.py @@ -1,4 +1,4 @@ -# Copyright 2024 Google LLC. +# Copyright 2025 Google LLC. # # 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/vmoe/evaluate/evaluator_test.py b/vmoe/evaluate/evaluator_test.py index 1147280..fb29cb4 100644 --- a/vmoe/evaluate/evaluator_test.py +++ b/vmoe/evaluate/evaluator_test.py @@ -1,4 +1,4 @@ -# Copyright 2024 Google LLC. +# Copyright 2025 Google LLC. # # 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/vmoe/evaluate/fewshot.py b/vmoe/evaluate/fewshot.py index 10b3e06..9d81fdd 100644 --- a/vmoe/evaluate/fewshot.py +++ b/vmoe/evaluate/fewshot.py @@ -1,4 +1,4 @@ -# Copyright 2024 Google LLC. +# Copyright 2025 Google LLC. # # 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/vmoe/evaluate/fewshot_test.py b/vmoe/evaluate/fewshot_test.py index a55ee9f..eb02eab 100644 --- a/vmoe/evaluate/fewshot_test.py +++ b/vmoe/evaluate/fewshot_test.py @@ -1,4 +1,4 @@ -# Copyright 2024 Google LLC. +# Copyright 2025 Google LLC. # # 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/vmoe/initialization/__init__.py b/vmoe/initialization/__init__.py index 7a1db00..6152b8f 100644 --- a/vmoe/initialization/__init__.py +++ b/vmoe/initialization/__init__.py @@ -1,4 +1,4 @@ -# Copyright 2024 Google LLC. +# Copyright 2025 Google LLC. # # 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/vmoe/initialization/initialization.py b/vmoe/initialization/initialization.py index 0bac15e..f99d7a5 100644 --- a/vmoe/initialization/initialization.py +++ b/vmoe/initialization/initialization.py @@ -1,4 +1,4 @@ -# Copyright 2024 Google LLC. +# Copyright 2025 Google LLC. # # 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/vmoe/initialization/initialization_test.py b/vmoe/initialization/initialization_test.py index 1b028a7..5a1da7a 100644 --- a/vmoe/initialization/initialization_test.py +++ b/vmoe/initialization/initialization_test.py @@ -1,4 +1,4 @@ -# Copyright 2024 Google LLC. +# Copyright 2025 Google LLC. # # 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/vmoe/initialization/mapping.py b/vmoe/initialization/mapping.py index 21239aa..447a4ce 100644 --- a/vmoe/initialization/mapping.py +++ b/vmoe/initialization/mapping.py @@ -1,4 +1,4 @@ -# Copyright 2024 Google LLC. +# Copyright 2025 Google LLC. # # 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/vmoe/initialization/mapping_test.py b/vmoe/initialization/mapping_test.py index 0b353e6..cf88851 100644 --- a/vmoe/initialization/mapping_test.py +++ b/vmoe/initialization/mapping_test.py @@ -1,4 +1,4 @@ -# Copyright 2024 Google LLC. +# Copyright 2025 Google LLC. # # 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/vmoe/initialization/rules.py b/vmoe/initialization/rules.py index e7d74f2..a6da9c6 100644 --- a/vmoe/initialization/rules.py +++ b/vmoe/initialization/rules.py @@ -1,4 +1,4 @@ -# Copyright 2024 Google LLC. +# Copyright 2025 Google LLC. # # 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/vmoe/initialization/rules_test.py b/vmoe/initialization/rules_test.py index e4a5a1f..7727a7f 100644 --- a/vmoe/initialization/rules_test.py +++ b/vmoe/initialization/rules_test.py @@ -1,4 +1,4 @@ -# Copyright 2024 Google LLC. +# Copyright 2025 Google LLC. # # 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/vmoe/moe.py b/vmoe/moe.py index 2992002..9815571 100644 --- a/vmoe/moe.py +++ b/vmoe/moe.py @@ -1,4 +1,4 @@ -# Copyright 2024 Google LLC. +# Copyright 2025 Google LLC. # # 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/vmoe/moe_test.py b/vmoe/moe_test.py index 8912847..3a7ec8d 100644 --- a/vmoe/moe_test.py +++ b/vmoe/moe_test.py @@ -1,4 +1,4 @@ -# Copyright 2024 Google LLC. +# Copyright 2025 Google LLC. # # 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/vmoe/multihost_utils.py b/vmoe/multihost_utils.py index 0474911..b5ad1d7 100644 --- a/vmoe/multihost_utils.py +++ b/vmoe/multihost_utils.py @@ -1,4 +1,4 @@ -# Copyright 2024 Google LLC. +# Copyright 2025 Google LLC. # # 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/vmoe/nn/ensemble_routing.py b/vmoe/nn/ensemble_routing.py index 786cce9..fd7c13f 100644 --- a/vmoe/nn/ensemble_routing.py +++ b/vmoe/nn/ensemble_routing.py @@ -1,4 +1,4 @@ -# Copyright 2024 Google LLC. +# Copyright 2025 Google LLC. # # 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/vmoe/nn/ensemble_routing_test.py b/vmoe/nn/ensemble_routing_test.py index a257374..dbecb55 100644 --- a/vmoe/nn/ensemble_routing_test.py +++ b/vmoe/nn/ensemble_routing_test.py @@ -1,4 +1,4 @@ -# Copyright 2024 Google LLC. +# Copyright 2025 Google LLC. # # 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/vmoe/nn/external.py b/vmoe/nn/external.py index 10edf12..bb3d22e 100644 --- a/vmoe/nn/external.py +++ b/vmoe/nn/external.py @@ -1,4 +1,4 @@ -# Copyright 2024 Google LLC. +# Copyright 2025 Google LLC. # # 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/vmoe/nn/external_test.py b/vmoe/nn/external_test.py index 25dc34f..70a9b1f 100644 --- a/vmoe/nn/external_test.py +++ b/vmoe/nn/external_test.py @@ -1,4 +1,4 @@ -# Copyright 2024 Google LLC. +# Copyright 2025 Google LLC. # # 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/vmoe/nn/models.py b/vmoe/nn/models.py index bba0807..0cb13ef 100644 --- a/vmoe/nn/models.py +++ b/vmoe/nn/models.py @@ -1,4 +1,4 @@ -# Copyright 2024 Google LLC. +# Copyright 2025 Google LLC. # # 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/vmoe/nn/routing.py b/vmoe/nn/routing.py index 0f7c311..dc86e1c 100644 --- a/vmoe/nn/routing.py +++ b/vmoe/nn/routing.py @@ -1,4 +1,4 @@ -# Copyright 2024 Google LLC. +# Copyright 2025 Google LLC. # # 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/vmoe/nn/routing_test.py b/vmoe/nn/routing_test.py index 570c570..461ba1f 100644 --- a/vmoe/nn/routing_test.py +++ b/vmoe/nn/routing_test.py @@ -1,4 +1,4 @@ -# Copyright 2024 Google LLC. +# Copyright 2025 Google LLC. # # 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/vmoe/nn/vit_moe.py b/vmoe/nn/vit_moe.py index 7319db4..ca6677b 100644 --- a/vmoe/nn/vit_moe.py +++ b/vmoe/nn/vit_moe.py @@ -1,4 +1,4 @@ -# Copyright 2024 Google LLC. +# Copyright 2025 Google LLC. # # 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/vmoe/nn/vit_moe_ensemble.py b/vmoe/nn/vit_moe_ensemble.py index b1f314d..30a6b2b 100644 --- a/vmoe/nn/vit_moe_ensemble.py +++ b/vmoe/nn/vit_moe_ensemble.py @@ -1,4 +1,4 @@ -# Copyright 2024 Google LLC. +# Copyright 2025 Google LLC. # # 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/vmoe/nn/vit_moe_ensemble_test.py b/vmoe/nn/vit_moe_ensemble_test.py index 73099d8..f198e7e 100644 --- a/vmoe/nn/vit_moe_ensemble_test.py +++ b/vmoe/nn/vit_moe_ensemble_test.py @@ -1,4 +1,4 @@ -# Copyright 2024 Google LLC. +# Copyright 2025 Google LLC. # # 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/vmoe/nn/vit_moe_test.py b/vmoe/nn/vit_moe_test.py index 2f42e08..7e7ed34 100644 --- a/vmoe/nn/vit_moe_test.py +++ b/vmoe/nn/vit_moe_test.py @@ -1,4 +1,4 @@ -# Copyright 2024 Google LLC. +# Copyright 2025 Google LLC. # # 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/vmoe/partitioning.py b/vmoe/partitioning.py index 01c15d1..7ac5327 100644 --- a/vmoe/partitioning.py +++ b/vmoe/partitioning.py @@ -1,4 +1,4 @@ -# Copyright 2024 Google LLC. +# Copyright 2025 Google LLC. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -87,7 +87,7 @@ def get_array_sharding_or_default(arr: jax.Array) -> jax.sharding.Sharding: if hasattr(arr, 'sharding'): return arr.sharding else: - return jax.sharding.GSPMDSharding.get_replicated(jax.devices()) + return NamedSharding(Mesh(jax.devices(), 'x'), PartitionSpec()) def process_has_contiguous_device_slice(devices: np.ndarray, diff --git a/vmoe/partitioning_test.py b/vmoe/partitioning_test.py index f531583..9d03440 100644 --- a/vmoe/partitioning_test.py +++ b/vmoe/partitioning_test.py @@ -1,4 +1,4 @@ -# Copyright 2024 Google LLC. +# Copyright 2025 Google LLC. # # 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/vmoe/projects/adversarial_attacks/attacks.py b/vmoe/projects/adversarial_attacks/attacks.py index ddacfa1..6488cbd 100644 --- a/vmoe/projects/adversarial_attacks/attacks.py +++ b/vmoe/projects/adversarial_attacks/attacks.py @@ -1,4 +1,4 @@ -# Copyright 2024 Google LLC. +# Copyright 2025 Google LLC. # # 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/vmoe/projects/adversarial_attacks/attacks_test.py b/vmoe/projects/adversarial_attacks/attacks_test.py index ee5a939..4a5e126 100644 --- a/vmoe/projects/adversarial_attacks/attacks_test.py +++ b/vmoe/projects/adversarial_attacks/attacks_test.py @@ -1,4 +1,4 @@ -# Copyright 2024 Google LLC. +# Copyright 2025 Google LLC. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -118,7 +118,7 @@ def loss_fn(x, y, metrics): return jax.numpy.square(x - y).sum() pgd_attack_stateless_jitted = jax.jit( - fun=functools.partial( + functools.partial( attacks.stateless_attack_pgd, max_epsilon=0.1, num_updates=1, apply_fn=apply_fn, loss_fn=loss_fn)) new_x, _ = pgd_attack_stateless_jitted(x, y, {}) diff --git a/vmoe/projects/adversarial_attacks/configs/attack/ilsvrc2012.py b/vmoe/projects/adversarial_attacks/configs/attack/ilsvrc2012.py index 69a722e..3b1ed3c 100644 --- a/vmoe/projects/adversarial_attacks/configs/attack/ilsvrc2012.py +++ b/vmoe/projects/adversarial_attacks/configs/attack/ilsvrc2012.py @@ -1,4 +1,4 @@ -# Copyright 2024 Google LLC. +# Copyright 2025 Google LLC. # # 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/vmoe/projects/adversarial_attacks/configs/common.py b/vmoe/projects/adversarial_attacks/configs/common.py index 24e2342..320ad82 100644 --- a/vmoe/projects/adversarial_attacks/configs/common.py +++ b/vmoe/projects/adversarial_attacks/configs/common.py @@ -1,4 +1,4 @@ -# Copyright 2024 Google LLC. +# Copyright 2025 Google LLC. # # 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/vmoe/projects/adversarial_attacks/lib.py b/vmoe/projects/adversarial_attacks/lib.py index 5a09740..b48ea3a 100644 --- a/vmoe/projects/adversarial_attacks/lib.py +++ b/vmoe/projects/adversarial_attacks/lib.py @@ -1,4 +1,4 @@ -# Copyright 2024 Google LLC. +# Copyright 2025 Google LLC. # # 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/vmoe/projects/adversarial_attacks/lib_test.py b/vmoe/projects/adversarial_attacks/lib_test.py index 9eacf6c..e8e82fa 100644 --- a/vmoe/projects/adversarial_attacks/lib_test.py +++ b/vmoe/projects/adversarial_attacks/lib_test.py @@ -1,4 +1,4 @@ -# Copyright 2024 Google LLC. +# Copyright 2025 Google LLC. # # 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/vmoe/projects/adversarial_attacks/main.py b/vmoe/projects/adversarial_attacks/main.py index 8cfd50b..15fec68 100644 --- a/vmoe/projects/adversarial_attacks/main.py +++ b/vmoe/projects/adversarial_attacks/main.py @@ -1,4 +1,4 @@ -# Copyright 2024 Google LLC. +# Copyright 2025 Google LLC. # # 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/vmoe/projects/adversarial_attacks/restore.py b/vmoe/projects/adversarial_attacks/restore.py index 5bf1245..b4dad81 100644 --- a/vmoe/projects/adversarial_attacks/restore.py +++ b/vmoe/projects/adversarial_attacks/restore.py @@ -1,4 +1,4 @@ -# Copyright 2024 Google LLC. +# Copyright 2025 Google LLC. # # 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/vmoe/projects/adversarial_attacks/restore_test.py b/vmoe/projects/adversarial_attacks/restore_test.py index 149de91..3042754 100644 --- a/vmoe/projects/adversarial_attacks/restore_test.py +++ b/vmoe/projects/adversarial_attacks/restore_test.py @@ -1,4 +1,4 @@ -# Copyright 2024 Google LLC. +# Copyright 2025 Google LLC. # # 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/vmoe/projects/contrastive/evaluators.py b/vmoe/projects/contrastive/evaluators.py index 7e0623a..fad5b64 100644 --- a/vmoe/projects/contrastive/evaluators.py +++ b/vmoe/projects/contrastive/evaluators.py @@ -1,4 +1,4 @@ -# Copyright 2024 Google LLC. +# Copyright 2025 Google LLC. # # 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/vmoe/projects/contrastive/models.py b/vmoe/projects/contrastive/models.py index bcf5764..5933c73 100644 --- a/vmoe/projects/contrastive/models.py +++ b/vmoe/projects/contrastive/models.py @@ -1,4 +1,4 @@ -# Copyright 2024 Google LLC. +# Copyright 2025 Google LLC. # # 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/vmoe/projects/contrastive/models_test.py b/vmoe/projects/contrastive/models_test.py index 3b59265..ee08df5 100644 --- a/vmoe/projects/contrastive/models_test.py +++ b/vmoe/projects/contrastive/models_test.py @@ -1,4 +1,4 @@ -# Copyright 2024 Google LLC. +# Copyright 2025 Google LLC. # # 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/vmoe/projects/contrastive/trainer.py b/vmoe/projects/contrastive/trainer.py index 99a088c..3347099 100644 --- a/vmoe/projects/contrastive/trainer.py +++ b/vmoe/projects/contrastive/trainer.py @@ -1,4 +1,4 @@ -# Copyright 2024 Google LLC. +# Copyright 2025 Google LLC. # # 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/vmoe/projects/soft_moe/configs/common.py b/vmoe/projects/soft_moe/configs/common.py index 9ba3560..a446d9b 100644 --- a/vmoe/projects/soft_moe/configs/common.py +++ b/vmoe/projects/soft_moe/configs/common.py @@ -1,4 +1,4 @@ -# Copyright 2024 Google LLC. +# Copyright 2025 Google LLC. # # 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/vmoe/projects/soft_moe/configs/pretrain_jft4b.py b/vmoe/projects/soft_moe/configs/pretrain_jft4b.py index c1551b6..2e667b5 100644 --- a/vmoe/projects/soft_moe/configs/pretrain_jft4b.py +++ b/vmoe/projects/soft_moe/configs/pretrain_jft4b.py @@ -1,4 +1,4 @@ -# Copyright 2024 Google LLC. +# Copyright 2025 Google LLC. # # 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/vmoe/projects/soft_moe/configs/pretrain_laion.py b/vmoe/projects/soft_moe/configs/pretrain_laion.py index bc3c643..2cab4f4 100644 --- a/vmoe/projects/soft_moe/configs/pretrain_laion.py +++ b/vmoe/projects/soft_moe/configs/pretrain_laion.py @@ -1,4 +1,4 @@ -# Copyright 2024 Google LLC. +# Copyright 2025 Google LLC. # # 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/vmoe/projects/soft_moe/main.py b/vmoe/projects/soft_moe/main.py index 652f7f3..aeada89 100644 --- a/vmoe/projects/soft_moe/main.py +++ b/vmoe/projects/soft_moe/main.py @@ -1,4 +1,4 @@ -# Copyright 2024 Google LLC. +# Copyright 2025 Google LLC. # # 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/vmoe/projects/soft_moe/main_contrastive.py b/vmoe/projects/soft_moe/main_contrastive.py index beffc62..1b3bb10 100644 --- a/vmoe/projects/soft_moe/main_contrastive.py +++ b/vmoe/projects/soft_moe/main_contrastive.py @@ -1,4 +1,4 @@ -# Copyright 2024 Google LLC. +# Copyright 2025 Google LLC. # # 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/vmoe/projects/soft_moe/router.py b/vmoe/projects/soft_moe/router.py index bff0978..16e71a1 100644 --- a/vmoe/projects/soft_moe/router.py +++ b/vmoe/projects/soft_moe/router.py @@ -1,4 +1,4 @@ -# Copyright 2024 Google LLC. +# Copyright 2025 Google LLC. # # 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/vmoe/projects/soft_moe/router_test.py b/vmoe/projects/soft_moe/router_test.py index f3b67ae..409db07 100644 --- a/vmoe/projects/soft_moe/router_test.py +++ b/vmoe/projects/soft_moe/router_test.py @@ -1,4 +1,4 @@ -# Copyright 2024 Google LLC. +# Copyright 2025 Google LLC. # # 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/vmoe/projects/sparsity_constrained_ot/kl_projection_routing.py b/vmoe/projects/sparsity_constrained_ot/kl_projection_routing.py index 8d3be4a..0efbc98 100644 --- a/vmoe/projects/sparsity_constrained_ot/kl_projection_routing.py +++ b/vmoe/projects/sparsity_constrained_ot/kl_projection_routing.py @@ -1,4 +1,4 @@ -# Copyright 2024 Google LLC. +# Copyright 2025 Google LLC. # # 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/vmoe/projects/sparsity_constrained_ot/kl_projection_routing_test.py b/vmoe/projects/sparsity_constrained_ot/kl_projection_routing_test.py index 661b6cc..33c2e27 100644 --- a/vmoe/projects/sparsity_constrained_ot/kl_projection_routing_test.py +++ b/vmoe/projects/sparsity_constrained_ot/kl_projection_routing_test.py @@ -1,4 +1,4 @@ -# Copyright 2024 Google LLC. +# Copyright 2025 Google LLC. # # 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/vmoe/projects/sparsity_constrained_ot/ksparse_projection_routing.py b/vmoe/projects/sparsity_constrained_ot/ksparse_projection_routing.py index cb8f336..ac50e18 100644 --- a/vmoe/projects/sparsity_constrained_ot/ksparse_projection_routing.py +++ b/vmoe/projects/sparsity_constrained_ot/ksparse_projection_routing.py @@ -1,4 +1,4 @@ -# Copyright 2024 Google LLC. +# Copyright 2025 Google LLC. # # 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/vmoe/projects/sparsity_constrained_ot/ksparse_projection_routing_test.py b/vmoe/projects/sparsity_constrained_ot/ksparse_projection_routing_test.py index fca49d1..ef0f088 100644 --- a/vmoe/projects/sparsity_constrained_ot/ksparse_projection_routing_test.py +++ b/vmoe/projects/sparsity_constrained_ot/ksparse_projection_routing_test.py @@ -1,4 +1,4 @@ -# Copyright 2024 Google LLC. +# Copyright 2025 Google LLC. # # 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/vmoe/projects/sparsity_constrained_ot/ot_routing.py b/vmoe/projects/sparsity_constrained_ot/ot_routing.py index 663b6f5..2043963 100644 --- a/vmoe/projects/sparsity_constrained_ot/ot_routing.py +++ b/vmoe/projects/sparsity_constrained_ot/ot_routing.py @@ -1,4 +1,4 @@ -# Copyright 2024 Google LLC. +# Copyright 2025 Google LLC. # # 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/vmoe/projects/sparsity_constrained_ot/sparse_projection_routing.py b/vmoe/projects/sparsity_constrained_ot/sparse_projection_routing.py index 879cde5..a488099 100644 --- a/vmoe/projects/sparsity_constrained_ot/sparse_projection_routing.py +++ b/vmoe/projects/sparsity_constrained_ot/sparse_projection_routing.py @@ -1,4 +1,4 @@ -# Copyright 2024 Google LLC. +# Copyright 2025 Google LLC. # # 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/vmoe/projects/sparsity_constrained_ot/sparse_projection_routing_test.py b/vmoe/projects/sparsity_constrained_ot/sparse_projection_routing_test.py index 4deeb82..518ce6b 100644 --- a/vmoe/projects/sparsity_constrained_ot/sparse_projection_routing_test.py +++ b/vmoe/projects/sparsity_constrained_ot/sparse_projection_routing_test.py @@ -1,4 +1,4 @@ -# Copyright 2024 Google LLC. +# Copyright 2025 Google LLC. # # 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/vmoe/train/main.py b/vmoe/train/main.py index 9274d13..b289b22 100644 --- a/vmoe/train/main.py +++ b/vmoe/train/main.py @@ -1,4 +1,4 @@ -# Copyright 2024 Google LLC. +# Copyright 2025 Google LLC. # # 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/vmoe/train/optimizer.py b/vmoe/train/optimizer.py index c569696..417d63d 100644 --- a/vmoe/train/optimizer.py +++ b/vmoe/train/optimizer.py @@ -1,4 +1,4 @@ -# Copyright 2024 Google LLC. +# Copyright 2025 Google LLC. # # 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/vmoe/train/optimizer_test.py b/vmoe/train/optimizer_test.py index df9b2da..b69382b 100644 --- a/vmoe/train/optimizer_test.py +++ b/vmoe/train/optimizer_test.py @@ -1,4 +1,4 @@ -# Copyright 2024 Google LLC. +# Copyright 2025 Google LLC. # # 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/vmoe/train/periodic_actions.py b/vmoe/train/periodic_actions.py index d2b4566..2b713fd 100644 --- a/vmoe/train/periodic_actions.py +++ b/vmoe/train/periodic_actions.py @@ -1,4 +1,4 @@ -# Copyright 2024 Google LLC. +# Copyright 2025 Google LLC. # # 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/vmoe/train/periodic_actions_test.py b/vmoe/train/periodic_actions_test.py index bbc2576..0e95b8b 100644 --- a/vmoe/train/periodic_actions_test.py +++ b/vmoe/train/periodic_actions_test.py @@ -1,4 +1,4 @@ -# Copyright 2024 Google LLC. +# Copyright 2025 Google LLC. # # 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/vmoe/train/schedule.py b/vmoe/train/schedule.py index 148766c..48a7725 100644 --- a/vmoe/train/schedule.py +++ b/vmoe/train/schedule.py @@ -1,4 +1,4 @@ -# Copyright 2024 Google LLC. +# Copyright 2025 Google LLC. # # 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/vmoe/train/schedule_test.py b/vmoe/train/schedule_test.py index 340480d..a30e74c 100644 --- a/vmoe/train/schedule_test.py +++ b/vmoe/train/schedule_test.py @@ -1,4 +1,4 @@ -# Copyright 2024 Google LLC. +# Copyright 2025 Google LLC. # # 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/vmoe/train/train_state.py b/vmoe/train/train_state.py index cd7627a..f733c8b 100644 --- a/vmoe/train/train_state.py +++ b/vmoe/train/train_state.py @@ -1,4 +1,4 @@ -# Copyright 2024 Google LLC. +# Copyright 2025 Google LLC. # # 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/vmoe/train/train_state_test.py b/vmoe/train/train_state_test.py index f6de995..e460feb 100644 --- a/vmoe/train/train_state_test.py +++ b/vmoe/train/train_state_test.py @@ -1,4 +1,4 @@ -# Copyright 2024 Google LLC. +# Copyright 2025 Google LLC. # # 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/vmoe/train/trainer.py b/vmoe/train/trainer.py index adcc652..4062a59 100644 --- a/vmoe/train/trainer.py +++ b/vmoe/train/trainer.py @@ -1,4 +1,4 @@ -# Copyright 2024 Google LLC. +# Copyright 2025 Google LLC. # # 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/vmoe/train/trainer_test.py b/vmoe/train/trainer_test.py index eab0b0f..6682994 100644 --- a/vmoe/train/trainer_test.py +++ b/vmoe/train/trainer_test.py @@ -1,4 +1,4 @@ -# Copyright 2024 Google LLC. +# Copyright 2025 Google LLC. # # 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/vmoe/train/tree_summarizer.py b/vmoe/train/tree_summarizer.py index 074e743..138c37a 100644 --- a/vmoe/train/tree_summarizer.py +++ b/vmoe/train/tree_summarizer.py @@ -1,4 +1,4 @@ -# Copyright 2024 Google LLC. +# Copyright 2025 Google LLC. # # 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/vmoe/train/tree_summarizer_test.py b/vmoe/train/tree_summarizer_test.py index ceb8ff6..de1aeb3 100644 --- a/vmoe/train/tree_summarizer_test.py +++ b/vmoe/train/tree_summarizer_test.py @@ -1,4 +1,4 @@ -# Copyright 2024 Google LLC. +# Copyright 2025 Google LLC. # # 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/vmoe/utils.py b/vmoe/utils.py index a097a9e..8c5f75c 100644 --- a/vmoe/utils.py +++ b/vmoe/utils.py @@ -1,4 +1,4 @@ -# Copyright 2024 Google LLC. +# Copyright 2025 Google LLC. # # 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/vmoe/utils_test.py b/vmoe/utils_test.py index ed8d0d8..cd4f1f2 100644 --- a/vmoe/utils_test.py +++ b/vmoe/utils_test.py @@ -1,4 +1,4 @@ -# Copyright 2024 Google LLC. +# Copyright 2025 Google LLC. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License.