diff --git a/hessian/model_debugger.py b/hessian/model_debugger.py index 9825653d..da294f86 100644 --- a/hessian/model_debugger.py +++ b/hessian/model_debugger.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/hessian/model_debugger_callback.py b/hessian/model_debugger_callback.py index e72f667c..a1b0485e 100644 --- a/hessian/model_debugger_callback.py +++ b/hessian/model_debugger_callback.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/hessian/precondition.py b/hessian/precondition.py index 6eba367c..77eefd11 100644 --- a/hessian/precondition.py +++ b/hessian/precondition.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/hessian/test_model_debugger.py b/hessian/test_model_debugger.py index 47f00204..2c6bd4ac 100644 --- a/hessian/test_model_debugger.py +++ b/hessian/test_model_debugger.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/hessian/test_precondition.py b/hessian/test_precondition.py index f8215a50..6cefd5a1 100644 --- a/hessian/test_precondition.py +++ b/hessian/test_precondition.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/__init__.py b/init2winit/__init__.py index 0e17992d..f1dbc852 100644 --- a/init2winit/__init__.py +++ b/init2winit/__init__.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/base_callback.py b/init2winit/base_callback.py index 711628a4..835c0f01 100644 --- a/init2winit/base_callback.py +++ b/init2winit/base_callback.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/callbacks.py b/init2winit/callbacks.py index 2d7f8c49..42940e00 100644 --- a/init2winit/callbacks.py +++ b/init2winit/callbacks.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/checkpoint.py b/init2winit/checkpoint.py index d1b02303..49ed335f 100644 --- a/init2winit/checkpoint.py +++ b/init2winit/checkpoint.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/dataset_lib/__init__.py b/init2winit/dataset_lib/__init__.py index 0e17992d..f1dbc852 100644 --- a/init2winit/dataset_lib/__init__.py +++ b/init2winit/dataset_lib/__init__.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/dataset_lib/autoaugment.py b/init2winit/dataset_lib/autoaugment.py index ec47bbcb..47bd45af 100644 --- a/init2winit/dataset_lib/autoaugment.py +++ b/init2winit/dataset_lib/autoaugment.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/dataset_lib/criteo_terabyte_dataset.py b/init2winit/dataset_lib/criteo_terabyte_dataset.py index 06a79bc8..db1268b8 100644 --- a/init2winit/dataset_lib/criteo_terabyte_dataset.py +++ b/init2winit/dataset_lib/criteo_terabyte_dataset.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/dataset_lib/data_selectors.py b/init2winit/dataset_lib/data_selectors.py index 5a9629e4..d64b0e0e 100644 --- a/init2winit/dataset_lib/data_selectors.py +++ b/init2winit/dataset_lib/data_selectors.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/dataset_lib/data_utils.py b/init2winit/dataset_lib/data_utils.py index 5f44be94..393bcb5f 100644 --- a/init2winit/dataset_lib/data_utils.py +++ b/init2winit/dataset_lib/data_utils.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/dataset_lib/datasets.py b/init2winit/dataset_lib/datasets.py index 5378a94f..ee27b326 100644 --- a/init2winit/dataset_lib/datasets.py +++ b/init2winit/dataset_lib/datasets.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit Authors. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -31,6 +31,7 @@ from init2winit.dataset_lib import nanodo_fineweb_edu from init2winit.dataset_lib import nqm_noise from init2winit.dataset_lib import ogbg_molpcba +from init2winit.dataset_lib import ogbg_molpcba_preprocessed from init2winit.dataset_lib import proteins from init2winit.dataset_lib import small_image_datasets from init2winit.dataset_lib import translate_wmt @@ -42,90 +43,145 @@ 'Dataset', ('getter', 'hparams', 'meta_data', 'fake_batch_getter')) _ALL_DATASETS = { - 'mnist': - _Dataset(small_image_datasets.get_mnist, - small_image_datasets.MNIST_HPARAMS, - small_image_datasets.MNIST_METADATA, None), - 'mnist_autoencoder': - _Dataset(small_image_datasets.get_mnist_autoencoder, - small_image_datasets.MNIST_AUTOENCODER_HPARAMS, - small_image_datasets.MNIST_AUTOENCODER_METADATA, None), - 'fashion_mnist': - _Dataset(small_image_datasets.get_fashion_mnist, - small_image_datasets.FASHION_MNIST_HPARAMS, - small_image_datasets.FASHION_MNIST_METADATA, None), - 'cifar10': - _Dataset(small_image_datasets.get_cifar10, - small_image_datasets.CIFAR10_DEFAULT_HPARAMS, - small_image_datasets.CIFAR10_METADATA, None), - 'cifar100': - _Dataset(small_image_datasets.get_cifar100, - small_image_datasets.CIFAR100_DEFAULT_HPARAMS, - small_image_datasets.CIFAR100_METADATA, None), - 'criteo1tb': - _Dataset(criteo_terabyte_dataset.get_criteo1tb, - criteo_terabyte_dataset.CRITEO1TB_DEFAULT_HPARAMS, - criteo_terabyte_dataset.CRITEO1TB_METADATA, - criteo_terabyte_dataset.get_fake_batch), - 'fake': - _Dataset(fake_dataset.get_fake, fake_dataset.DEFAULT_HPARAMS, - fake_dataset.METADATA, fake_dataset.get_fake_batch), - 'fastmri': - _Dataset(fastmri_dataset.get_fastmri, fastmri_dataset.DEFAULT_HPARAMS, - fastmri_dataset.METADATA, fastmri_dataset.get_fake_batch), + 'mnist': _Dataset( + small_image_datasets.get_mnist, + small_image_datasets.MNIST_HPARAMS, + small_image_datasets.MNIST_METADATA, + None, + ), + 'mnist_autoencoder': _Dataset( + small_image_datasets.get_mnist_autoencoder, + small_image_datasets.MNIST_AUTOENCODER_HPARAMS, + small_image_datasets.MNIST_AUTOENCODER_METADATA, + None, + ), + 'fashion_mnist': _Dataset( + small_image_datasets.get_fashion_mnist, + small_image_datasets.FASHION_MNIST_HPARAMS, + small_image_datasets.FASHION_MNIST_METADATA, + None, + ), + 'cifar10': _Dataset( + small_image_datasets.get_cifar10, + small_image_datasets.CIFAR10_DEFAULT_HPARAMS, + small_image_datasets.CIFAR10_METADATA, + None, + ), + 'cifar100': _Dataset( + small_image_datasets.get_cifar100, + small_image_datasets.CIFAR100_DEFAULT_HPARAMS, + small_image_datasets.CIFAR100_METADATA, + None, + ), + 'criteo1tb': _Dataset( + criteo_terabyte_dataset.get_criteo1tb, + criteo_terabyte_dataset.CRITEO1TB_DEFAULT_HPARAMS, + criteo_terabyte_dataset.CRITEO1TB_METADATA, + criteo_terabyte_dataset.get_fake_batch, + ), + 'fake': _Dataset( + fake_dataset.get_fake, + fake_dataset.DEFAULT_HPARAMS, + fake_dataset.METADATA, + fake_dataset.get_fake_batch, + ), + 'fastmri': _Dataset( + fastmri_dataset.get_fastmri, + fastmri_dataset.DEFAULT_HPARAMS, + fastmri_dataset.METADATA, + fastmri_dataset.get_fake_batch, + ), 'fineweb_edu_10B': _Dataset( fineweb_edu_10b.get_fineweb_edu, fineweb_edu_10b.DEFAULT_HPARAMS, - fineweb_edu_10b.METADATA, None), - 'imagenet': - _Dataset(imagenet_dataset.get_imagenet, - imagenet_dataset.DEFAULT_HPARAMS, imagenet_dataset.METADATA, - imagenet_dataset.get_fake_batch), - 'translate_wmt': - _Dataset(translate_wmt.get_translate_wmt, translate_wmt.DEFAULT_HPARAMS, - translate_wmt.METADATA, translate_wmt.get_fake_batch), - 'librispeech': - _Dataset(librispeech.get_librispeech, librispeech.DEFAULT_HPARAMS, - librispeech.METADATA, librispeech.get_fake_batch), - 'lm1b_v2': - _Dataset(lm1b_v2.get_lm1b, lm1b_v2.DEFAULT_HPARAMS, lm1b_v2.METADATA, - None), - 'mlperf_imagenet': - _Dataset(mlperf_imagenet_dataset.get_mlperf_imagenet, - mlperf_imagenet_dataset.DEFAULT_HPARAMS, - mlperf_imagenet_dataset.METADATA, - mlperf_imagenet_dataset.get_fake_batch), - 'svhn_no_extra': - _Dataset(small_image_datasets.get_svhn_no_extra, - small_image_datasets.SVHN_NO_EXTRA_DEFAULT_HPARAMS, - small_image_datasets.SVHN_NO_EXTRA_METADATA, None), + fineweb_edu_10b.METADATA, + None, + ), + 'imagenet': _Dataset( + imagenet_dataset.get_imagenet, + imagenet_dataset.DEFAULT_HPARAMS, + imagenet_dataset.METADATA, + imagenet_dataset.get_fake_batch, + ), + 'translate_wmt': _Dataset( + translate_wmt.get_translate_wmt, + translate_wmt.DEFAULT_HPARAMS, + translate_wmt.METADATA, + translate_wmt.get_fake_batch, + ), + 'librispeech': _Dataset( + librispeech.get_librispeech, + librispeech.DEFAULT_HPARAMS, + librispeech.METADATA, + librispeech.get_fake_batch, + ), + 'lm1b_v2': _Dataset( + lm1b_v2.get_lm1b, lm1b_v2.DEFAULT_HPARAMS, lm1b_v2.METADATA, None + ), + 'mlperf_imagenet': _Dataset( + mlperf_imagenet_dataset.get_mlperf_imagenet, + mlperf_imagenet_dataset.DEFAULT_HPARAMS, + mlperf_imagenet_dataset.METADATA, + mlperf_imagenet_dataset.get_fake_batch, + ), + 'svhn_no_extra': _Dataset( + small_image_datasets.get_svhn_no_extra, + small_image_datasets.SVHN_NO_EXTRA_DEFAULT_HPARAMS, + small_image_datasets.SVHN_NO_EXTRA_METADATA, + None, + ), 'c4': _Dataset( nanodo_c4.get_dataset, nanodo_c4.DEFAULT_HPARAMS, - nanodo_c4.METADATA, None), + nanodo_c4.METADATA, + None, + ), 'fineweb_edu': _Dataset( nanodo_fineweb_edu.get_dataset, nanodo_fineweb_edu.DEFAULT_HPARAMS, - nanodo_fineweb_edu.METADATA, None), - 'nqm_noise': - _Dataset(nqm_noise.get_nqm_noise, nqm_noise.NQM_HPARAMS, - nqm_noise.NQM_METADATA, None), - 'ogbg_molpcba': - _Dataset(ogbg_molpcba.get_ogbg_molpcba, ogbg_molpcba.DEFAULT_HPARAMS, - ogbg_molpcba.METADATA, ogbg_molpcba.get_fake_batch), - 'uniref50': - _Dataset(proteins.get_uniref, proteins.DEFAULT_HPARAMS, - proteins.METADATA, None), - 'wikitext2': - _Dataset(wikitext2.get_wikitext2, wikitext2.DEFAULT_HPARAMS, - wikitext2.METADATA, None), - 'wikitext103': - _Dataset(wikitext103.get_wikitext103, wikitext103.DEFAULT_HPARAMS, - wikitext2.METADATA, None), - 'wikitext103_spm': - _Dataset(wikitext103_spm.get_wikitext103, - wikitext103_spm.DEFAULT_HPARAMS, - wikitext103_spm.METADATA, None), + nanodo_fineweb_edu.METADATA, + None, + ), + 'nqm_noise': _Dataset( + nqm_noise.get_nqm_noise, + nqm_noise.NQM_HPARAMS, + nqm_noise.NQM_METADATA, + None, + ), + 'ogbg_molpcba': _Dataset( + ogbg_molpcba.get_ogbg_molpcba, + ogbg_molpcba.DEFAULT_HPARAMS, + ogbg_molpcba.METADATA, + ogbg_molpcba.get_fake_batch, + ), + 'ogbg_molpcba_preprocessed': _Dataset( + ogbg_molpcba_preprocessed.get_ogbg_molpcba_preprocessed, + ogbg_molpcba_preprocessed.DEFAULT_HPARAMS, + ogbg_molpcba_preprocessed.METADATA, + # Reuse logic for fake batch if needed, or None + ogbg_molpcba.get_fake_batch, + ), + 'uniref50': _Dataset( + proteins.get_uniref, proteins.DEFAULT_HPARAMS, proteins.METADATA, None + ), + 'wikitext2': _Dataset( + wikitext2.get_wikitext2, + wikitext2.DEFAULT_HPARAMS, + wikitext2.METADATA, + None, + ), + 'wikitext103': _Dataset( + wikitext103.get_wikitext103, + wikitext103.DEFAULT_HPARAMS, + wikitext2.METADATA, + None, + ), + 'wikitext103_spm': _Dataset( + wikitext103_spm.get_wikitext103, + wikitext103_spm.DEFAULT_HPARAMS, + wikitext103_spm.METADATA, + None, + ), } @@ -238,4 +294,3 @@ def get_data_selector(selector_name: Optional[str]): 'Unrecognized selector: {}'.format(selector_name)) from None return selector - diff --git a/init2winit/dataset_lib/fake_dataset.py b/init2winit/dataset_lib/fake_dataset.py index 059156aa..332b6cb8 100644 --- a/init2winit/dataset_lib/fake_dataset.py +++ b/init2winit/dataset_lib/fake_dataset.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/dataset_lib/fastmri_dataset.py b/init2winit/dataset_lib/fastmri_dataset.py index c3196d7e..bed2f6ef 100644 --- a/init2winit/dataset_lib/fastmri_dataset.py +++ b/init2winit/dataset_lib/fastmri_dataset.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/dataset_lib/fineweb_edu_10b.py b/init2winit/dataset_lib/fineweb_edu_10b.py index 71efb1be..de38b69c 100644 --- a/init2winit/dataset_lib/fineweb_edu_10b.py +++ b/init2winit/dataset_lib/fineweb_edu_10b.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/dataset_lib/fineweb_edu_10b_input_pipeline.py b/init2winit/dataset_lib/fineweb_edu_10b_input_pipeline.py index 0000051b..10f0ccfd 100644 --- a/init2winit/dataset_lib/fineweb_edu_10b_input_pipeline.py +++ b/init2winit/dataset_lib/fineweb_edu_10b_input_pipeline.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/dataset_lib/image_preprocessing.py b/init2winit/dataset_lib/image_preprocessing.py index 449c5f57..7c674fd3 100644 --- a/init2winit/dataset_lib/image_preprocessing.py +++ b/init2winit/dataset_lib/image_preprocessing.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/dataset_lib/imagenet_dataset.py b/init2winit/dataset_lib/imagenet_dataset.py index e0ab4c6d..49c4852b 100644 --- a/init2winit/dataset_lib/imagenet_dataset.py +++ b/init2winit/dataset_lib/imagenet_dataset.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/dataset_lib/imagenet_preprocessing.py b/init2winit/dataset_lib/imagenet_preprocessing.py index 35b91b77..c5bfdd17 100644 --- a/init2winit/dataset_lib/imagenet_preprocessing.py +++ b/init2winit/dataset_lib/imagenet_preprocessing.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/dataset_lib/librispeech.py b/init2winit/dataset_lib/librispeech.py index 9360fefa..d79f9e3b 100644 --- a/init2winit/dataset_lib/librispeech.py +++ b/init2winit/dataset_lib/librispeech.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/dataset_lib/librispeech_input_pipeline.py b/init2winit/dataset_lib/librispeech_input_pipeline.py index 6ef7aeba..d9d9cdca 100644 --- a/init2winit/dataset_lib/librispeech_input_pipeline.py +++ b/init2winit/dataset_lib/librispeech_input_pipeline.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/dataset_lib/lm1b_input_pipeline_v2.py b/init2winit/dataset_lib/lm1b_input_pipeline_v2.py index 895412e2..8494ab8f 100644 --- a/init2winit/dataset_lib/lm1b_input_pipeline_v2.py +++ b/init2winit/dataset_lib/lm1b_input_pipeline_v2.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/dataset_lib/lm1b_v2.py b/init2winit/dataset_lib/lm1b_v2.py index 917b4d19..b2c718c9 100644 --- a/init2winit/dataset_lib/lm1b_v2.py +++ b/init2winit/dataset_lib/lm1b_v2.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/dataset_lib/mlperf_imagenet_dataset.py b/init2winit/dataset_lib/mlperf_imagenet_dataset.py index 00c5056e..59b85276 100644 --- a/init2winit/dataset_lib/mlperf_imagenet_dataset.py +++ b/init2winit/dataset_lib/mlperf_imagenet_dataset.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/dataset_lib/mlperf_input_pipeline.py b/init2winit/dataset_lib/mlperf_input_pipeline.py index 11a16563..ca4c3c13 100644 --- a/init2winit/dataset_lib/mlperf_input_pipeline.py +++ b/init2winit/dataset_lib/mlperf_input_pipeline.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/dataset_lib/mt_pipeline.py b/init2winit/dataset_lib/mt_pipeline.py index 7318d3a7..bfc37c2d 100644 --- a/init2winit/dataset_lib/mt_pipeline.py +++ b/init2winit/dataset_lib/mt_pipeline.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/dataset_lib/mt_pipeline_test.py b/init2winit/dataset_lib/mt_pipeline_test.py index f41ade72..0eb42fbf 100644 --- a/init2winit/dataset_lib/mt_pipeline_test.py +++ b/init2winit/dataset_lib/mt_pipeline_test.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/dataset_lib/mt_tokenizer.py b/init2winit/dataset_lib/mt_tokenizer.py index f5d86e54..3efaa257 100644 --- a/init2winit/dataset_lib/mt_tokenizer.py +++ b/init2winit/dataset_lib/mt_tokenizer.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/dataset_lib/nanodo_c4.py b/init2winit/dataset_lib/nanodo_c4.py index bd537ecf..170a74b2 100644 --- a/init2winit/dataset_lib/nanodo_c4.py +++ b/init2winit/dataset_lib/nanodo_c4.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/dataset_lib/nanodo_data_loader_shared.py b/init2winit/dataset_lib/nanodo_data_loader_shared.py index 8d551b8b..d3b057da 100644 --- a/init2winit/dataset_lib/nanodo_data_loader_shared.py +++ b/init2winit/dataset_lib/nanodo_data_loader_shared.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/dataset_lib/nanodo_fineweb_edu.py b/init2winit/dataset_lib/nanodo_fineweb_edu.py index d43b6b1d..5876c5df 100644 --- a/init2winit/dataset_lib/nanodo_fineweb_edu.py +++ b/init2winit/dataset_lib/nanodo_fineweb_edu.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/dataset_lib/nqm_noise.py b/init2winit/dataset_lib/nqm_noise.py index b8dbcacf..e5642da3 100644 --- a/init2winit/dataset_lib/nqm_noise.py +++ b/init2winit/dataset_lib/nqm_noise.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/dataset_lib/ogbg_molpcba.py b/init2winit/dataset_lib/ogbg_molpcba.py index 176a60b4..90d11bf2 100644 --- a/init2winit/dataset_lib/ogbg_molpcba.py +++ b/init2winit/dataset_lib/ogbg_molpcba.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit Authors. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -285,10 +285,20 @@ def _get_batch_iterator(dataset_iter, weights_shards = [] -def get_ogbg_molpcba(shuffle_rng, batch_size, eval_batch_size, hps=None): +def get_ogbg_molpcba( + shuffle_rng, + batch_size, + eval_batch_size, + hps=None, + override_process_count=None, +): """Data generators for ogbg-molpcba.""" - process_count = jax.process_count() + process_count = ( + jax.process_count() + if not override_process_count + else override_process_count + ) if batch_size % process_count != 0: raise ValueError( 'process_count={} must divide batch_size={}.'.format( @@ -339,7 +349,9 @@ def get_ogbg_molpcba(shuffle_rng, batch_size, eval_batch_size, hps=None): edges_per_graph=int(max_edges_multiplier), add_bidirectional_edges=hps.add_bidirectional_edges, add_virtual_node=hps.add_virtual_node, - add_self_loops=hps.add_self_loops) + add_self_loops=hps.add_self_loops, + num_shards=override_process_count, + ) def train_iterator_fn(): return iterator_from_ds( diff --git a/init2winit/dataset_lib/ogbg_molpcba_preprocessed.py b/init2winit/dataset_lib/ogbg_molpcba_preprocessed.py new file mode 100644 index 00000000..fb0027a8 --- /dev/null +++ b/init2winit/dataset_lib/ogbg_molpcba_preprocessed.py @@ -0,0 +1,142 @@ +# coding=utf-8 +# Copyright 2026 The init2winit Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Data generator for preprocessed ogbg_molpcba.""" + +from init2winit.dataset_lib import data_utils +from init2winit.dataset_lib import ogbg_molpcba +import jax +from ml_collections.config_dict import config_dict +import tensorflow as tf + +# Reuse default hparams from ogbg_molpcba, but add dataset_path +DEFAULT_HPARAMS = config_dict.ConfigDict(ogbg_molpcba.DEFAULT_HPARAMS) +DEFAULT_HPARAMS.dataset_path = '' + + +METADATA = ogbg_molpcba.METADATA + + +def get_ogbg_molpcba_preprocessed( + shuffle_rng, batch_size, eval_batch_size, hps=None +): + """Get preprocessed ogbg_molpcba dataset. + + The saved batches must have the correct per-host batch size (batch_size / + process_count). + + Args: + shuffle_rng: RNG key for shuffling (used for fallback). + batch_size: Global batch size. + eval_batch_size: Global eval batch size. + hps: Hyperparameters. + + Returns: + Dataset object. + """ + if not hps.dataset_path: + raise ValueError( + 'hps.dataset_path must be provided for ogbg_molpcba_preprocessed' + ) + + # Fallback to original dataset for validation/test splits + original_ds = ogbg_molpcba.get_ogbg_molpcba( + shuffle_rng, + batch_size, + eval_batch_size, + hps, + override_process_count=None, + ) + + def train_iterator_fn(): + ds = tf.data.Dataset.load(hps.dataset_path, compression='GZIP') + + process_count = jax.process_count() + process_index = jax.process_index() + + if process_count > 1: + ds = ds.shard(process_count, process_index) + + iterator = iter(ds.as_numpy_iterator()) + return iterator + + return data_utils.Dataset( + train_iterator_fn, + original_ds.eval_train_epoch, + original_ds.valid_epoch, + original_ds.test_epoch, + ) + + +def generate_and_save_dataset( + hps, + output_path, + num_steps, + batch_size, + seed=0, + progress_bar_fn=None, + target_num_processes=None, +): + """Generates and saves the ogbg_molpcba dataset. + + Args: + hps: Hyperparameters. + output_path: Path to save the dataset. + num_steps: Number of steps (batches) to generate. + batch_size: Batch size. + seed: Random seed. + progress_bar_fn: Optional function to wrap the iterator for progress + tracking. Should accept an iterable and return an iterable. + target_num_processes: If set, the dataset will be sharded for this number of + processes (=number of TPUs usually). + """ + rng = jax.random.PRNGKey(seed) + + dataset_builder = ogbg_molpcba.get_ogbg_molpcba( + shuffle_rng=rng, + batch_size=batch_size, + eval_batch_size=batch_size, + hps=hps, + override_process_count=target_num_processes, + ) + + train_iter = dataset_builder.train_iterator_fn() + + def generator(): + iterator = range(num_steps) + if progress_bar_fn: + iterator = progress_bar_fn(iterator) + + for _ in iterator: + try: + batch = next(train_iter) + yield batch + except StopIteration: + return + + # Determine output signature from the first batch + peek_iter = dataset_builder.train_iterator_fn() + first_batch = next(peek_iter) + + output_signature = tf.nest.map_structure( + lambda x: tf.TensorSpec(shape=x.shape, dtype=tf.as_dtype(x.dtype)), + first_batch, + ) + + ds = tf.data.Dataset.from_generator( + generator, output_signature=output_signature + ) + + ds.save(output_path, compression='GZIP') diff --git a/init2winit/dataset_lib/pg19.py b/init2winit/dataset_lib/pg19.py index 2ee8c412..db2afec5 100644 --- a/init2winit/dataset_lib/pg19.py +++ b/init2winit/dataset_lib/pg19.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/dataset_lib/protein_vocab.py b/init2winit/dataset_lib/protein_vocab.py index 88375d72..3a2de57a 100644 --- a/init2winit/dataset_lib/protein_vocab.py +++ b/init2winit/dataset_lib/protein_vocab.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/dataset_lib/proteins.py b/init2winit/dataset_lib/proteins.py index 7b647c43..347f0b38 100644 --- a/init2winit/dataset_lib/proteins.py +++ b/init2winit/dataset_lib/proteins.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/dataset_lib/small_image_datasets.py b/init2winit/dataset_lib/small_image_datasets.py index 82c3c553..34202128 100644 --- a/init2winit/dataset_lib/small_image_datasets.py +++ b/init2winit/dataset_lib/small_image_datasets.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/dataset_lib/spm_tokenizer.py b/init2winit/dataset_lib/spm_tokenizer.py index 993874c3..49b3cbe3 100644 --- a/init2winit/dataset_lib/spm_tokenizer.py +++ b/init2winit/dataset_lib/spm_tokenizer.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/dataset_lib/test_data_utils.py b/init2winit/dataset_lib/test_data_utils.py index 27a408ae..184a0287 100644 --- a/init2winit/dataset_lib/test_data_utils.py +++ b/init2winit/dataset_lib/test_data_utils.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/dataset_lib/test_datasets.py b/init2winit/dataset_lib/test_datasets.py index 921f9547..1e9866a7 100644 --- a/init2winit/dataset_lib/test_datasets.py +++ b/init2winit/dataset_lib/test_datasets.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/dataset_lib/test_datasets_multihost.py b/init2winit/dataset_lib/test_datasets_multihost.py index 601c93c6..9de322e0 100644 --- a/init2winit/dataset_lib/test_datasets_multihost.py +++ b/init2winit/dataset_lib/test_datasets_multihost.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/dataset_lib/test_fineweb_edu_10b_input_pipeline.py b/init2winit/dataset_lib/test_fineweb_edu_10b_input_pipeline.py index f6a2a6d1..97ab105b 100644 --- a/init2winit/dataset_lib/test_fineweb_edu_10b_input_pipeline.py +++ b/init2winit/dataset_lib/test_fineweb_edu_10b_input_pipeline.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/dataset_lib/test_ogbg_molpcba.py b/init2winit/dataset_lib/test_ogbg_molpcba.py index 9fc3bbfd..659e3765 100644 --- a/init2winit/dataset_lib/test_ogbg_molpcba.py +++ b/init2winit/dataset_lib/test_ogbg_molpcba.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/dataset_lib/test_ogbg_molpcba_preprocessed.py b/init2winit/dataset_lib/test_ogbg_molpcba_preprocessed.py new file mode 100644 index 00000000..b12d572a --- /dev/null +++ b/init2winit/dataset_lib/test_ogbg_molpcba_preprocessed.py @@ -0,0 +1,149 @@ +# coding=utf-8 +# Copyright 2026 The init2winit Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for ogbg_molpcba_preprocessed.""" + +import logging +import os +import shutil +import tempfile +from unittest import mock + +from absl.testing import absltest +from init2winit.dataset_lib import datasets +from init2winit.dataset_lib import ogbg_molpcba +from init2winit.dataset_lib import ogbg_molpcba_preprocessed +import jax +import jraph +from ml_collections import config_dict +import numpy as np + +logging.basicConfig( + level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s' +) + + +class OgbgMolpcbaPreprocessedTest(absltest.TestCase): + """Tests for ogbg_molpcba_preprocessed. + + This test mocks ogbg_molpcba.get_ogbg_molpcba to return a dummy dataset + iterator, and checks that generate_and_save_dataset and the + ogbg_molpcba_preprocessed dataset loader work as expected by verifying that + the data can be loaded and iterated over, and that the data matches the + original dummy dataset. + """ + + def setUp(self): + super().setUp() + self.test_dir = tempfile.mkdtemp() + self.dataset_path = os.path.join(self.test_dir, 'test_dataset') + + # Create dummy batch for mocking + self.dummy_batch = { + 'inputs': jraph.GraphsTuple( + nodes=np.random.randn(2, 2).astype(np.float32), + edges=np.random.randn(2, 2).astype(np.float32), + receivers=np.array([0, 1], dtype=np.int32), + senders=np.array([0, 1], dtype=np.int32), + n_node=np.array([2], dtype=np.int32), + n_edge=np.array([2], dtype=np.int32), + globals=np.zeros((1, 2), dtype=np.float32), + ), + 'targets': np.random.randn(1, 128).astype(np.float32), + 'weights': np.random.randn(1, 128).astype(np.float32), + } + + # Mock ogbg_molpcba.get_ogbg_molpcba to return a dummy builder + self.patcher = mock.patch.object(ogbg_molpcba, 'get_ogbg_molpcba') + self.mock_get_ogbg = self.patcher.start() + + mock_ds = mock.Mock() + # Return iterator that yields dummy_batch infinitely (or enough times) + mock_ds.train_iterator_fn.return_value = iter([self.dummy_batch] * 100) + self.mock_get_ogbg.return_value = mock_ds + + # Generate a small dataset + hps = config_dict.ConfigDict(ogbg_molpcba.DEFAULT_HPARAMS) + logging.info('Generating dataset...') + ogbg_molpcba_preprocessed.generate_and_save_dataset( + hps=hps, + output_path=self.dataset_path, + num_steps=5, + batch_size=2, + seed=0, + ) + logging.info('Dataset generated at %s', self.dataset_path) + + def tearDown(self): + self.patcher.stop() + shutil.rmtree(self.test_dir) + super().tearDown() + + def test_load_dataset_train_iterator(self): + """Tests that the preprocessed dataset can be loaded and iterated over.""" + + hps = config_dict.ConfigDict(ogbg_molpcba.DEFAULT_HPARAMS) + hps.dataset_path = self.dataset_path + + # Call the loader directly or via datasets registry + logging.info('Loading dataset...') + loader = datasets.get_dataset('ogbg_molpcba_preprocessed') + ds = loader( + shuffle_rng=jax.random.PRNGKey(0), + batch_size=2, + eval_batch_size=2, + hps=hps, + ) + + # Verify train iterator + train_iter = ds.train_iterator_fn() + batch = next(train_iter) + self.assertTrue(hasattr(batch['inputs'], 'nodes')) + self.assertTrue(hasattr(batch['inputs'], 'edges')) + self.assertTrue(hasattr(batch['inputs'], 'senders')) + self.assertTrue(hasattr(batch['inputs'], 'receivers')) + self.assertNotEmpty(batch['inputs'].nodes) + + def assert_arrays_equal(arr1, arr2): + np.testing.assert_allclose(arr1, arr2) + + assert_arrays_equal(batch['inputs'].nodes, self.dummy_batch['inputs'].nodes) + assert_arrays_equal(batch['inputs'].edges, self.dummy_batch['inputs'].edges) + assert_arrays_equal( + batch['inputs'].senders, self.dummy_batch['inputs'].senders + ) + assert_arrays_equal( + batch['inputs'].receivers, self.dummy_batch['inputs'].receivers + ) + assert_arrays_equal( + batch['inputs'].globals, self.dummy_batch['inputs'].globals + ) + assert_arrays_equal( + batch['inputs'].n_node, self.dummy_batch['inputs'].n_node + ) + assert_arrays_equal( + batch['inputs'].n_edge, self.dummy_batch['inputs'].n_edge + ) + + assert_arrays_equal(batch['targets'], self.dummy_batch['targets']) + assert_arrays_equal(batch['weights'], self.dummy_batch['weights']) + + # Verify that we can iterate multiple times (it loads from disk) + batch2 = next(train_iter) + self.assertTrue(hasattr(batch2['inputs'], 'nodes')) + + +if __name__ == '__main__': + absltest.main() diff --git a/init2winit/dataset_lib/test_small_image_datasets.py b/init2winit/dataset_lib/test_small_image_datasets.py index c970ee3d..74b9a2ca 100644 --- a/init2winit/dataset_lib/test_small_image_datasets.py +++ b/init2winit/dataset_lib/test_small_image_datasets.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/dataset_lib/test_wikitext_tokenizer.py b/init2winit/dataset_lib/test_wikitext_tokenizer.py index 16691f39..f289c49e 100644 --- a/init2winit/dataset_lib/test_wikitext_tokenizer.py +++ b/init2winit/dataset_lib/test_wikitext_tokenizer.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/dataset_lib/translate_wmt.py b/init2winit/dataset_lib/translate_wmt.py index 84c75b4a..f6e7331d 100644 --- a/init2winit/dataset_lib/translate_wmt.py +++ b/init2winit/dataset_lib/translate_wmt.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/dataset_lib/wikitext103.py b/init2winit/dataset_lib/wikitext103.py index 94de9182..440826e8 100644 --- a/init2winit/dataset_lib/wikitext103.py +++ b/init2winit/dataset_lib/wikitext103.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/dataset_lib/wikitext103_input_pipeline.py b/init2winit/dataset_lib/wikitext103_input_pipeline.py index b79ea38e..7e7f1167 100644 --- a/init2winit/dataset_lib/wikitext103_input_pipeline.py +++ b/init2winit/dataset_lib/wikitext103_input_pipeline.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/dataset_lib/wikitext103_spm.py b/init2winit/dataset_lib/wikitext103_spm.py index 34a09535..c9fbf6db 100644 --- a/init2winit/dataset_lib/wikitext103_spm.py +++ b/init2winit/dataset_lib/wikitext103_spm.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/dataset_lib/wikitext2.py b/init2winit/dataset_lib/wikitext2.py index acb0c5ab..2eb71b04 100644 --- a/init2winit/dataset_lib/wikitext2.py +++ b/init2winit/dataset_lib/wikitext2.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/dataset_lib/wikitext2_input_pipeline.py b/init2winit/dataset_lib/wikitext2_input_pipeline.py index 667caf70..79cab2ef 100644 --- a/init2winit/dataset_lib/wikitext2_input_pipeline.py +++ b/init2winit/dataset_lib/wikitext2_input_pipeline.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/dataset_lib/wikitext_tokenizer.py b/init2winit/dataset_lib/wikitext_tokenizer.py index 168ea3a2..6425fa48 100644 --- a/init2winit/dataset_lib/wikitext_tokenizer.py +++ b/init2winit/dataset_lib/wikitext_tokenizer.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/gradient_statistics_callback.py b/init2winit/gradient_statistics_callback.py index af123cec..044dc4be 100644 --- a/init2winit/gradient_statistics_callback.py +++ b/init2winit/gradient_statistics_callback.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/hyperparameters.py b/init2winit/hyperparameters.py index c708b688..a366391d 100644 --- a/init2winit/hyperparameters.py +++ b/init2winit/hyperparameters.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/init_lib/__init__.py b/init2winit/init_lib/__init__.py index 0e17992d..f1dbc852 100644 --- a/init2winit/init_lib/__init__.py +++ b/init2winit/init_lib/__init__.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/init_lib/initializers.py b/init2winit/init_lib/initializers.py index b1f664de..feb277f7 100644 --- a/init2winit/init_lib/initializers.py +++ b/init2winit/init_lib/initializers.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/init_lib/meta_init.py b/init2winit/init_lib/meta_init.py index 9ed6be62..400dfb13 100644 --- a/init2winit/init_lib/meta_init.py +++ b/init2winit/init_lib/meta_init.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/init_lib/sparse_init.py b/init2winit/init_lib/sparse_init.py index 17a98b8c..2a9caef3 100644 --- a/init2winit/init_lib/sparse_init.py +++ b/init2winit/init_lib/sparse_init.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/init_lib/test_initializers.py b/init2winit/init_lib/test_initializers.py index 15c4fd7a..ef836fb6 100644 --- a/init2winit/init_lib/test_initializers.py +++ b/init2winit/init_lib/test_initializers.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/main.py b/init2winit/main.py index 5d959d36..69fc74ec 100644 --- a/init2winit/main.py +++ b/init2winit/main.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/main_config_flags.py b/init2winit/main_config_flags.py index b7b061e4..ef991754 100644 --- a/init2winit/main_config_flags.py +++ b/init2winit/main_config_flags.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/model_lib/__init__.py b/init2winit/model_lib/__init__.py index 0e17992d..f1dbc852 100644 --- a/init2winit/model_lib/__init__.py +++ b/init2winit/model_lib/__init__.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/model_lib/adabelief_densenet.py b/init2winit/model_lib/adabelief_densenet.py index 0785dab3..af64b028 100644 --- a/init2winit/model_lib/adabelief_densenet.py +++ b/init2winit/model_lib/adabelief_densenet.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/model_lib/adabelief_resnet.py b/init2winit/model_lib/adabelief_resnet.py index 580d4ec7..72d6282b 100644 --- a/init2winit/model_lib/adabelief_resnet.py +++ b/init2winit/model_lib/adabelief_resnet.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/model_lib/adabelief_vgg.py b/init2winit/model_lib/adabelief_vgg.py index 2aaa4710..d7e1b143 100644 --- a/init2winit/model_lib/adabelief_vgg.py +++ b/init2winit/model_lib/adabelief_vgg.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/model_lib/attention.py b/init2winit/model_lib/attention.py index 7befc5b0..0d919f27 100644 --- a/init2winit/model_lib/attention.py +++ b/init2winit/model_lib/attention.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/model_lib/autoencoder.py b/init2winit/model_lib/autoencoder.py index ebf53751..36c803d7 100644 --- a/init2winit/model_lib/autoencoder.py +++ b/init2winit/model_lib/autoencoder.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/model_lib/base_model.py b/init2winit/model_lib/base_model.py index be089c62..ca04ee5c 100644 --- a/init2winit/model_lib/base_model.py +++ b/init2winit/model_lib/base_model.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit Authors. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -202,8 +202,7 @@ def initialize(self, initializer, hps, rng, metrics_logger): start_time = time.time() model_init_fn = jax.jit( - functools.partial(self.flax_module.init, train=False), - backend='cpu') + functools.partial(self.flax_module.init, train=False)) init_dict = model_init_fn({'params': params_rng, 'dropout': dropout_rng}, *fake_input_batch) diff --git a/init2winit/model_lib/binarize_layers.py b/init2winit/model_lib/binarize_layers.py index 5709acc6..b953473c 100644 --- a/init2winit/model_lib/binarize_layers.py +++ b/init2winit/model_lib/binarize_layers.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/model_lib/conformer.py b/init2winit/model_lib/conformer.py index 64c1f64d..ccfa5f0f 100644 --- a/init2winit/model_lib/conformer.py +++ b/init2winit/model_lib/conformer.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/model_lib/convolutional_autoencoder.py b/init2winit/model_lib/convolutional_autoencoder.py index 2a584268..f04fc459 100644 --- a/init2winit/model_lib/convolutional_autoencoder.py +++ b/init2winit/model_lib/convolutional_autoencoder.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/model_lib/deepspeech.py b/init2winit/model_lib/deepspeech.py index d0e32567..2d489e8e 100644 --- a/init2winit/model_lib/deepspeech.py +++ b/init2winit/model_lib/deepspeech.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/model_lib/dlrm.py b/init2winit/model_lib/dlrm.py index 4df724cb..b340ab56 100644 --- a/init2winit/model_lib/dlrm.py +++ b/init2winit/model_lib/dlrm.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/model_lib/fully_connected.py b/init2winit/model_lib/fully_connected.py index 3c6249f5..641e057e 100644 --- a/init2winit/model_lib/fully_connected.py +++ b/init2winit/model_lib/fully_connected.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/model_lib/gnn.py b/init2winit/model_lib/gnn.py index 42f42d3e..33ff062e 100644 --- a/init2winit/model_lib/gnn.py +++ b/init2winit/model_lib/gnn.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/model_lib/librispeech_preprocessor.py b/init2winit/model_lib/librispeech_preprocessor.py index 2f08cf8b..6ff1ce21 100644 --- a/init2winit/model_lib/librispeech_preprocessor.py +++ b/init2winit/model_lib/librispeech_preprocessor.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/model_lib/local_attention_transformer.py b/init2winit/model_lib/local_attention_transformer.py index 2ea78801..e9bf08a6 100644 --- a/init2winit/model_lib/local_attention_transformer.py +++ b/init2winit/model_lib/local_attention_transformer.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/model_lib/losses.py b/init2winit/model_lib/losses.py index 02793cb7..9dfc5635 100644 --- a/init2winit/model_lib/losses.py +++ b/init2winit/model_lib/losses.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/model_lib/lstm.py b/init2winit/model_lib/lstm.py index 50294e1b..96225b60 100644 --- a/init2winit/model_lib/lstm.py +++ b/init2winit/model_lib/lstm.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/model_lib/lstm_lm.py b/init2winit/model_lib/lstm_lm.py index cee9a09d..af0c333d 100644 --- a/init2winit/model_lib/lstm_lm.py +++ b/init2winit/model_lib/lstm_lm.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/model_lib/max_pooling_cnn.py b/init2winit/model_lib/max_pooling_cnn.py index 6abb4f39..1520c7a9 100644 --- a/init2winit/model_lib/max_pooling_cnn.py +++ b/init2winit/model_lib/max_pooling_cnn.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/model_lib/metrics.py b/init2winit/model_lib/metrics.py index 5f404ffe..90de68d6 100644 --- a/init2winit/model_lib/metrics.py +++ b/init2winit/model_lib/metrics.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/model_lib/metrics_minimized_registry.py b/init2winit/model_lib/metrics_minimized_registry.py index be5d0f04..cc15b5ed 100644 --- a/init2winit/model_lib/metrics_minimized_registry.py +++ b/init2winit/model_lib/metrics_minimized_registry.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/model_lib/mlperf_resnet.py b/init2winit/model_lib/mlperf_resnet.py index 8d2e8027..59c8d562 100644 --- a/init2winit/model_lib/mlperf_resnet.py +++ b/init2winit/model_lib/mlperf_resnet.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/model_lib/model_utils.py b/init2winit/model_lib/model_utils.py index bfb6742f..beeb1492 100644 --- a/init2winit/model_lib/model_utils.py +++ b/init2winit/model_lib/model_utils.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/model_lib/models.py b/init2winit/model_lib/models.py index edeb3280..38bc1c1e 100644 --- a/init2winit/model_lib/models.py +++ b/init2winit/model_lib/models.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/model_lib/nanodo.py b/init2winit/model_lib/nanodo.py index 6812ca4b..99fb0a0e 100644 --- a/init2winit/model_lib/nanodo.py +++ b/init2winit/model_lib/nanodo.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/model_lib/normalization.py b/init2winit/model_lib/normalization.py index 2f22a6ea..d57f67f8 100644 --- a/init2winit/model_lib/normalization.py +++ b/init2winit/model_lib/normalization.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/model_lib/nqm.py b/init2winit/model_lib/nqm.py index 7502b628..9c51ffe4 100644 --- a/init2winit/model_lib/nqm.py +++ b/init2winit/model_lib/nqm.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/model_lib/partition_tree.py b/init2winit/model_lib/partition_tree.py index 5891d3fc..80b05ee0 100644 --- a/init2winit/model_lib/partition_tree.py +++ b/init2winit/model_lib/partition_tree.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/model_lib/resnet.py b/init2winit/model_lib/resnet.py index 36bea7bb..ae312232 100644 --- a/init2winit/model_lib/resnet.py +++ b/init2winit/model_lib/resnet.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/model_lib/rope_nanodo.py b/init2winit/model_lib/rope_nanodo.py index 8b57eb56..12210690 100644 --- a/init2winit/model_lib/rope_nanodo.py +++ b/init2winit/model_lib/rope_nanodo.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/model_lib/simple_cnn.py b/init2winit/model_lib/simple_cnn.py index 2af0d370..9cbdb5e6 100644 --- a/init2winit/model_lib/simple_cnn.py +++ b/init2winit/model_lib/simple_cnn.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/model_lib/spectrum_augmenter.py b/init2winit/model_lib/spectrum_augmenter.py index b4b0d206..69f20277 100644 --- a/init2winit/model_lib/spectrum_augmenter.py +++ b/init2winit/model_lib/spectrum_augmenter.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/model_lib/test_local_attention_transformer.py b/init2winit/model_lib/test_local_attention_transformer.py index e542cbf7..aa8c2949 100644 --- a/init2winit/model_lib/test_local_attention_transformer.py +++ b/init2winit/model_lib/test_local_attention_transformer.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/model_lib/test_losses.py b/init2winit/model_lib/test_losses.py index b81e246c..e32573d5 100644 --- a/init2winit/model_lib/test_losses.py +++ b/init2winit/model_lib/test_losses.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/model_lib/test_metrics.py b/init2winit/model_lib/test_metrics.py index ff92fa6a..622b8012 100644 --- a/init2winit/model_lib/test_metrics.py +++ b/init2winit/model_lib/test_metrics.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/model_lib/test_models.py b/init2winit/model_lib/test_models.py index 42987737..fd46f514 100644 --- a/init2winit/model_lib/test_models.py +++ b/init2winit/model_lib/test_models.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/model_lib/test_normalization.py b/init2winit/model_lib/test_normalization.py index b82f40c1..75f0d33f 100644 --- a/init2winit/model_lib/test_normalization.py +++ b/init2winit/model_lib/test_normalization.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/model_lib/transformer_lm.py b/init2winit/model_lib/transformer_lm.py index f30b7d87..0e310ff8 100644 --- a/init2winit/model_lib/transformer_lm.py +++ b/init2winit/model_lib/transformer_lm.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/model_lib/transformer_stu_lm.py b/init2winit/model_lib/transformer_stu_lm.py index 5742d51c..07e41acf 100644 --- a/init2winit/model_lib/transformer_stu_lm.py +++ b/init2winit/model_lib/transformer_stu_lm.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/model_lib/transformer_stu_tensordot_lm.py b/init2winit/model_lib/transformer_stu_tensordot_lm.py index 4ae0e007..42a1ae52 100644 --- a/init2winit/model_lib/transformer_stu_tensordot_lm.py +++ b/init2winit/model_lib/transformer_stu_tensordot_lm.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/model_lib/unet.py b/init2winit/model_lib/unet.py index 622a73ae..3a4e3741 100644 --- a/init2winit/model_lib/unet.py +++ b/init2winit/model_lib/unet.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/model_lib/vit.py b/init2winit/model_lib/vit.py index 7ad6b81a..6bcad8f9 100644 --- a/init2winit/model_lib/vit.py +++ b/init2winit/model_lib/vit.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/model_lib/wide_resnet.py b/init2winit/model_lib/wide_resnet.py index a1ff6eae..8df134a9 100644 --- a/init2winit/model_lib/wide_resnet.py +++ b/init2winit/model_lib/wide_resnet.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/model_lib/xformer_translate.py b/init2winit/model_lib/xformer_translate.py index 312ed046..bc85c054 100644 --- a/init2winit/model_lib/xformer_translate.py +++ b/init2winit/model_lib/xformer_translate.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/model_lib/xformer_translate_binary.py b/init2winit/model_lib/xformer_translate_binary.py index 52906d04..66082507 100644 --- a/init2winit/model_lib/xformer_translate_binary.py +++ b/init2winit/model_lib/xformer_translate_binary.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/model_lib/xformer_translate_mlc_variant.py b/init2winit/model_lib/xformer_translate_mlc_variant.py index 68386552..f90d36d5 100644 --- a/init2winit/model_lib/xformer_translate_mlc_variant.py +++ b/init2winit/model_lib/xformer_translate_mlc_variant.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/mt_eval/decode.py b/init2winit/mt_eval/decode.py index 8cd91d02..34531674 100644 --- a/init2winit/mt_eval/decode.py +++ b/init2winit/mt_eval/decode.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/mt_eval/eval_utils.py b/init2winit/mt_eval/eval_utils.py index 9c44b29c..427d221d 100644 --- a/init2winit/mt_eval/eval_utils.py +++ b/init2winit/mt_eval/eval_utils.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/mt_eval/inference.py b/init2winit/mt_eval/inference.py index d83bcda0..aa148bc8 100644 --- a/init2winit/mt_eval/inference.py +++ b/init2winit/mt_eval/inference.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/mt_eval/main.py b/init2winit/mt_eval/main.py index ce0fe033..aeac0d7c 100644 --- a/init2winit/mt_eval/main.py +++ b/init2winit/mt_eval/main.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/mt_eval/mt_callback.py b/init2winit/mt_eval/mt_callback.py index 09a2de85..fd0cfdbe 100644 --- a/init2winit/mt_eval/mt_callback.py +++ b/init2winit/mt_eval/mt_callback.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/optimizer_lib/__init__.py b/init2winit/optimizer_lib/__init__.py index 0e17992d..f1dbc852 100644 --- a/init2winit/optimizer_lib/__init__.py +++ b/init2winit/optimizer_lib/__init__.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/optimizer_lib/factor_sam.py b/init2winit/optimizer_lib/factor_sam.py index 4e4d04fe..a2d668fa 100644 --- a/init2winit/optimizer_lib/factor_sam.py +++ b/init2winit/optimizer_lib/factor_sam.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/optimizer_lib/gradient_accumulator.py b/init2winit/optimizer_lib/gradient_accumulator.py index 16ceac3b..7dbe9005 100644 --- a/init2winit/optimizer_lib/gradient_accumulator.py +++ b/init2winit/optimizer_lib/gradient_accumulator.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/optimizer_lib/kitchen_sink/__init__.py b/init2winit/optimizer_lib/kitchen_sink/__init__.py index c4503d03..92a44a83 100644 --- a/init2winit/optimizer_lib/kitchen_sink/__init__.py +++ b/init2winit/optimizer_lib/kitchen_sink/__init__.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/optimizer_lib/kitchen_sink/_src/alias.py b/init2winit/optimizer_lib/kitchen_sink/_src/alias.py index 1ed4efa4..347b2f32 100644 --- a/init2winit/optimizer_lib/kitchen_sink/_src/alias.py +++ b/init2winit/optimizer_lib/kitchen_sink/_src/alias.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/optimizer_lib/kitchen_sink/_src/combine.py b/init2winit/optimizer_lib/kitchen_sink/_src/combine.py index fbc42ed8..6164e4e6 100644 --- a/init2winit/optimizer_lib/kitchen_sink/_src/combine.py +++ b/init2winit/optimizer_lib/kitchen_sink/_src/combine.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/optimizer_lib/kitchen_sink/_src/core.py b/init2winit/optimizer_lib/kitchen_sink/_src/core.py index ba62df89..634a8ea1 100644 --- a/init2winit/optimizer_lib/kitchen_sink/_src/core.py +++ b/init2winit/optimizer_lib/kitchen_sink/_src/core.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/optimizer_lib/kitchen_sink/_src/mask.py b/init2winit/optimizer_lib/kitchen_sink/_src/mask.py index 34994af8..f89f352e 100644 --- a/init2winit/optimizer_lib/kitchen_sink/_src/mask.py +++ b/init2winit/optimizer_lib/kitchen_sink/_src/mask.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/optimizer_lib/kitchen_sink/_src/preconditioner.py b/init2winit/optimizer_lib/kitchen_sink/_src/preconditioner.py index b826a1e1..5778f259 100644 --- a/init2winit/optimizer_lib/kitchen_sink/_src/preconditioner.py +++ b/init2winit/optimizer_lib/kitchen_sink/_src/preconditioner.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/optimizer_lib/kitchen_sink/_src/test_core.py b/init2winit/optimizer_lib/kitchen_sink/_src/test_core.py index 3aee8e76..74cd7f53 100644 --- a/init2winit/optimizer_lib/kitchen_sink/_src/test_core.py +++ b/init2winit/optimizer_lib/kitchen_sink/_src/test_core.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/optimizer_lib/kitchen_sink/_src/test_mask.py b/init2winit/optimizer_lib/kitchen_sink/_src/test_mask.py index 12390b4e..310ed91f 100644 --- a/init2winit/optimizer_lib/kitchen_sink/_src/test_mask.py +++ b/init2winit/optimizer_lib/kitchen_sink/_src/test_mask.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/optimizer_lib/kitchen_sink/_src/test_preconditioner.py b/init2winit/optimizer_lib/kitchen_sink/_src/test_preconditioner.py index 4073e03d..842e3f8a 100644 --- a/init2winit/optimizer_lib/kitchen_sink/_src/test_preconditioner.py +++ b/init2winit/optimizer_lib/kitchen_sink/_src/test_preconditioner.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/optimizer_lib/kitchen_sink/_src/test_transform.py b/init2winit/optimizer_lib/kitchen_sink/_src/test_transform.py index 597579f4..796fc4b4 100644 --- a/init2winit/optimizer_lib/kitchen_sink/_src/test_transform.py +++ b/init2winit/optimizer_lib/kitchen_sink/_src/test_transform.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/optimizer_lib/kitchen_sink/_src/transform.py b/init2winit/optimizer_lib/kitchen_sink/_src/transform.py index 194a8c21..db147e14 100644 --- a/init2winit/optimizer_lib/kitchen_sink/_src/transform.py +++ b/init2winit/optimizer_lib/kitchen_sink/_src/transform.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/optimizer_lib/kitchen_sink/_src/utils.py b/init2winit/optimizer_lib/kitchen_sink/_src/utils.py index 073154b2..77503971 100644 --- a/init2winit/optimizer_lib/kitchen_sink/_src/utils.py +++ b/init2winit/optimizer_lib/kitchen_sink/_src/utils.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/optimizer_lib/linalg/low_rank_root_update.py b/init2winit/optimizer_lib/linalg/low_rank_root_update.py index 0366dcdf..ece83085 100644 --- a/init2winit/optimizer_lib/linalg/low_rank_root_update.py +++ b/init2winit/optimizer_lib/linalg/low_rank_root_update.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/optimizer_lib/linalg/low_rank_root_update_test.py b/init2winit/optimizer_lib/linalg/low_rank_root_update_test.py index d82fd461..500dd171 100644 --- a/init2winit/optimizer_lib/linalg/low_rank_root_update_test.py +++ b/init2winit/optimizer_lib/linalg/low_rank_root_update_test.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/optimizer_lib/linalg/paterson_stockmeyer.py b/init2winit/optimizer_lib/linalg/paterson_stockmeyer.py index b5e1c9bb..310f38d9 100644 --- a/init2winit/optimizer_lib/linalg/paterson_stockmeyer.py +++ b/init2winit/optimizer_lib/linalg/paterson_stockmeyer.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/optimizer_lib/linalg/pth_inv_root_rmn.py b/init2winit/optimizer_lib/linalg/pth_inv_root_rmn.py index 0579dd47..72049772 100644 --- a/init2winit/optimizer_lib/linalg/pth_inv_root_rmn.py +++ b/init2winit/optimizer_lib/linalg/pth_inv_root_rmn.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/optimizer_lib/linalg/pth_inv_root_rmn_coefficients.py b/init2winit/optimizer_lib/linalg/pth_inv_root_rmn_coefficients.py index 0f77fbd9..cf3b5f3c 100644 --- a/init2winit/optimizer_lib/linalg/pth_inv_root_rmn_coefficients.py +++ b/init2winit/optimizer_lib/linalg/pth_inv_root_rmn_coefficients.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/optimizer_lib/linalg/pth_inv_root_rmn_test.py b/init2winit/optimizer_lib/linalg/pth_inv_root_rmn_test.py index db55dc0f..16978e28 100644 --- a/init2winit/optimizer_lib/linalg/pth_inv_root_rmn_test.py +++ b/init2winit/optimizer_lib/linalg/pth_inv_root_rmn_test.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/optimizer_lib/linalg/root_selector.py b/init2winit/optimizer_lib/linalg/root_selector.py index 63ce96cd..2a4dd170 100644 --- a/init2winit/optimizer_lib/linalg/root_selector.py +++ b/init2winit/optimizer_lib/linalg/root_selector.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/optimizer_lib/muon.py b/init2winit/optimizer_lib/muon.py index cfff4f87..cc799ab0 100644 --- a/init2winit/optimizer_lib/muon.py +++ b/init2winit/optimizer_lib/muon.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/optimizer_lib/online_newton_step.py b/init2winit/optimizer_lib/online_newton_step.py index 7c84dad8..2bd495a9 100644 --- a/init2winit/optimizer_lib/online_newton_step.py +++ b/init2winit/optimizer_lib/online_newton_step.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/optimizer_lib/optimizers.py b/init2winit/optimizer_lib/optimizers.py index 1c1780d0..a20344d8 100644 --- a/init2winit/optimizer_lib/optimizers.py +++ b/init2winit/optimizer_lib/optimizers.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/optimizer_lib/pax_adafactor.py b/init2winit/optimizer_lib/pax_adafactor.py index c2762b43..53671a87 100644 --- a/init2winit/optimizer_lib/pax_adafactor.py +++ b/init2winit/optimizer_lib/pax_adafactor.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/optimizer_lib/samuel.py b/init2winit/optimizer_lib/samuel.py index fa01e30f..5b6a0f1f 100644 --- a/init2winit/optimizer_lib/samuel.py +++ b/init2winit/optimizer_lib/samuel.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/optimizer_lib/search_subspace.py b/init2winit/optimizer_lib/search_subspace.py index 1b9aaa5c..58bc33a1 100644 --- a/init2winit/optimizer_lib/search_subspace.py +++ b/init2winit/optimizer_lib/search_subspace.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/optimizer_lib/sharpness_aware_minimization.py b/init2winit/optimizer_lib/sharpness_aware_minimization.py index 97fb1314..2d16924d 100644 --- a/init2winit/optimizer_lib/sharpness_aware_minimization.py +++ b/init2winit/optimizer_lib/sharpness_aware_minimization.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/optimizer_lib/sla.py b/init2winit/optimizer_lib/sla.py index 70ed1ffe..bc9df7db 100644 --- a/init2winit/optimizer_lib/sla.py +++ b/init2winit/optimizer_lib/sla.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/optimizer_lib/sla_test.py b/init2winit/optimizer_lib/sla_test.py index 7d759859..1af02664 100644 --- a/init2winit/optimizer_lib/sla_test.py +++ b/init2winit/optimizer_lib/sla_test.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/optimizer_lib/test_gradient_accumulator.py b/init2winit/optimizer_lib/test_gradient_accumulator.py index f46f8fe4..123b5837 100644 --- a/init2winit/optimizer_lib/test_gradient_accumulator.py +++ b/init2winit/optimizer_lib/test_gradient_accumulator.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/optimizer_lib/test_optimizers.py b/init2winit/optimizer_lib/test_optimizers.py index a0d03e28..7a3484fb 100644 --- a/init2winit/optimizer_lib/test_optimizers.py +++ b/init2winit/optimizer_lib/test_optimizers.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/optimizer_lib/test_search_subspace.py b/init2winit/optimizer_lib/test_search_subspace.py index 263ebdc7..49b640b5 100644 --- a/init2winit/optimizer_lib/test_search_subspace.py +++ b/init2winit/optimizer_lib/test_search_subspace.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/optimizer_lib/test_utils.py b/init2winit/optimizer_lib/test_utils.py index 7d2825b1..0e231840 100644 --- a/init2winit/optimizer_lib/test_utils.py +++ b/init2winit/optimizer_lib/test_utils.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/optimizer_lib/utils.py b/init2winit/optimizer_lib/utils.py index e54ded6f..edbb6aa5 100644 --- a/init2winit/optimizer_lib/utils.py +++ b/init2winit/optimizer_lib/utils.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/projects/optlrschedule/notebook_utils/pandas_util.py b/init2winit/projects/optlrschedule/notebook_utils/pandas_util.py index d60f1514..110f8ada 100644 --- a/init2winit/projects/optlrschedule/notebook_utils/pandas_util.py +++ b/init2winit/projects/optlrschedule/notebook_utils/pandas_util.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/projects/optlrschedule/notebook_utils/parquet_util.py b/init2winit/projects/optlrschedule/notebook_utils/parquet_util.py index 32556feb..121e4fa1 100644 --- a/init2winit/projects/optlrschedule/notebook_utils/parquet_util.py +++ b/init2winit/projects/optlrschedule/notebook_utils/parquet_util.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/projects/optlrschedule/notebook_utils/plot_util.py b/init2winit/projects/optlrschedule/notebook_utils/plot_util.py index 96101c2c..4de7aa8c 100644 --- a/init2winit/projects/optlrschedule/notebook_utils/plot_util.py +++ b/init2winit/projects/optlrschedule/notebook_utils/plot_util.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/projects/optlrschedule/notebook_utils/schedule_util.py b/init2winit/projects/optlrschedule/notebook_utils/schedule_util.py index 1a207073..3c13c2a2 100644 --- a/init2winit/projects/optlrschedule/notebook_utils/schedule_util.py +++ b/init2winit/projects/optlrschedule/notebook_utils/schedule_util.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/projects/optlrschedule/notebook_utils/test_pandas_util.py b/init2winit/projects/optlrschedule/notebook_utils/test_pandas_util.py index b29e6a94..8c746df9 100644 --- a/init2winit/projects/optlrschedule/notebook_utils/test_pandas_util.py +++ b/init2winit/projects/optlrschedule/notebook_utils/test_pandas_util.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/projects/optlrschedule/run_search_decoupled.py b/init2winit/projects/optlrschedule/run_search_decoupled.py index e3c04944..2a9fff4c 100644 --- a/init2winit/projects/optlrschedule/run_search_decoupled.py +++ b/init2winit/projects/optlrschedule/run_search_decoupled.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/projects/optlrschedule/scheduler/base_schedule_family.py b/init2winit/projects/optlrschedule/scheduler/base_schedule_family.py index 3f612177..67f4bb4e 100644 --- a/init2winit/projects/optlrschedule/scheduler/base_schedule_family.py +++ b/init2winit/projects/optlrschedule/scheduler/base_schedule_family.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/projects/optlrschedule/scheduler/constant_schedule_family.py b/init2winit/projects/optlrschedule/scheduler/constant_schedule_family.py index e8e31331..646d8237 100644 --- a/init2winit/projects/optlrschedule/scheduler/constant_schedule_family.py +++ b/init2winit/projects/optlrschedule/scheduler/constant_schedule_family.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/projects/optlrschedule/scheduler/cosine_schedule_family.py b/init2winit/projects/optlrschedule/scheduler/cosine_schedule_family.py index 2ecbe9cf..05ff3568 100644 --- a/init2winit/projects/optlrschedule/scheduler/cosine_schedule_family.py +++ b/init2winit/projects/optlrschedule/scheduler/cosine_schedule_family.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/projects/optlrschedule/scheduler/cosine_standard_schedule_family.py b/init2winit/projects/optlrschedule/scheduler/cosine_standard_schedule_family.py index 559725cf..d74ae1e9 100644 --- a/init2winit/projects/optlrschedule/scheduler/cosine_standard_schedule_family.py +++ b/init2winit/projects/optlrschedule/scheduler/cosine_standard_schedule_family.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/projects/optlrschedule/scheduler/cosine_y_schedule_family.py b/init2winit/projects/optlrschedule/scheduler/cosine_y_schedule_family.py index 20856e61..8526a930 100644 --- a/init2winit/projects/optlrschedule/scheduler/cosine_y_schedule_family.py +++ b/init2winit/projects/optlrschedule/scheduler/cosine_y_schedule_family.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/projects/optlrschedule/scheduler/piecewise_schedule_family.py b/init2winit/projects/optlrschedule/scheduler/piecewise_schedule_family.py index 5ed86a26..b713cb5a 100644 --- a/init2winit/projects/optlrschedule/scheduler/piecewise_schedule_family.py +++ b/init2winit/projects/optlrschedule/scheduler/piecewise_schedule_family.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/projects/optlrschedule/scheduler/rex_schedule_family.py b/init2winit/projects/optlrschedule/scheduler/rex_schedule_family.py index f5c8d113..c0e50f97 100644 --- a/init2winit/projects/optlrschedule/scheduler/rex_schedule_family.py +++ b/init2winit/projects/optlrschedule/scheduler/rex_schedule_family.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/projects/optlrschedule/scheduler/schedule_families.py b/init2winit/projects/optlrschedule/scheduler/schedule_families.py index 67b61f69..14e09c33 100644 --- a/init2winit/projects/optlrschedule/scheduler/schedule_families.py +++ b/init2winit/projects/optlrschedule/scheduler/schedule_families.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/projects/optlrschedule/scheduler/smooth_nonmonotonic_schedule_family.py b/init2winit/projects/optlrschedule/scheduler/smooth_nonmonotonic_schedule_family.py index db640f07..237eac8f 100644 --- a/init2winit/projects/optlrschedule/scheduler/smooth_nonmonotonic_schedule_family.py +++ b/init2winit/projects/optlrschedule/scheduler/smooth_nonmonotonic_schedule_family.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/projects/optlrschedule/scheduler/sqrt_schedule_family.py b/init2winit/projects/optlrschedule/scheduler/sqrt_schedule_family.py index 62f7c6f7..710ed4bb 100644 --- a/init2winit/projects/optlrschedule/scheduler/sqrt_schedule_family.py +++ b/init2winit/projects/optlrschedule/scheduler/sqrt_schedule_family.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/projects/optlrschedule/scheduler/test_scheduler.py b/init2winit/projects/optlrschedule/scheduler/test_scheduler.py index 3eecd445..e9f6dda7 100644 --- a/init2winit/projects/optlrschedule/scheduler/test_scheduler.py +++ b/init2winit/projects/optlrschedule/scheduler/test_scheduler.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/projects/optlrschedule/scheduler/twopointslinear_schedule_family.py b/init2winit/projects/optlrschedule/scheduler/twopointslinear_schedule_family.py index c3e40598..5ca70553 100644 --- a/init2winit/projects/optlrschedule/scheduler/twopointslinear_schedule_family.py +++ b/init2winit/projects/optlrschedule/scheduler/twopointslinear_schedule_family.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/projects/optlrschedule/scheduler/twopointsspline_schedule_family.py b/init2winit/projects/optlrschedule/scheduler/twopointsspline_schedule_family.py index 900ecfc0..feb01676 100644 --- a/init2winit/projects/optlrschedule/scheduler/twopointsspline_schedule_family.py +++ b/init2winit/projects/optlrschedule/scheduler/twopointsspline_schedule_family.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/projects/optlrschedule/scheduler/twopointsspline_y_schedule_family.py b/init2winit/projects/optlrschedule/scheduler/twopointsspline_y_schedule_family.py index 3ca05b11..7d19a6b7 100644 --- a/init2winit/projects/optlrschedule/scheduler/twopointsspline_y_schedule_family.py +++ b/init2winit/projects/optlrschedule/scheduler/twopointsspline_y_schedule_family.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/projects/optlrschedule/search_algorithm/coordinate_descent_search.py b/init2winit/projects/optlrschedule/search_algorithm/coordinate_descent_search.py index f62ca6d5..21f522ac 100644 --- a/init2winit/projects/optlrschedule/search_algorithm/coordinate_descent_search.py +++ b/init2winit/projects/optlrschedule/search_algorithm/coordinate_descent_search.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/projects/optlrschedule/search_algorithm/grid_search.py b/init2winit/projects/optlrschedule/search_algorithm/grid_search.py index aeb1ef9b..b199e34f 100644 --- a/init2winit/projects/optlrschedule/search_algorithm/grid_search.py +++ b/init2winit/projects/optlrschedule/search_algorithm/grid_search.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/projects/optlrschedule/search_algorithm/random_search.py b/init2winit/projects/optlrschedule/search_algorithm/random_search.py index 757f043b..950f0ffa 100644 --- a/init2winit/projects/optlrschedule/search_algorithm/random_search.py +++ b/init2winit/projects/optlrschedule/search_algorithm/random_search.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/projects/optlrschedule/search_algorithm/search_algorithms.py b/init2winit/projects/optlrschedule/search_algorithm/search_algorithms.py index 01d5d08b..39fa93a0 100644 --- a/init2winit/projects/optlrschedule/search_algorithm/search_algorithms.py +++ b/init2winit/projects/optlrschedule/search_algorithm/search_algorithms.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/projects/optlrschedule/search_algorithm/test_search_algorithm.py b/init2winit/projects/optlrschedule/search_algorithm/test_search_algorithm.py index 54fb3925..5d9c7858 100644 --- a/init2winit/projects/optlrschedule/search_algorithm/test_search_algorithm.py +++ b/init2winit/projects/optlrschedule/search_algorithm/test_search_algorithm.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/projects/optlrschedule/workload/base_workload.py b/init2winit/projects/optlrschedule/workload/base_workload.py index b2d75a43..5152c8b7 100644 --- a/init2winit/projects/optlrschedule/workload/base_workload.py +++ b/init2winit/projects/optlrschedule/workload/base_workload.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/projects/optlrschedule/workload/cifar10_cnn.py b/init2winit/projects/optlrschedule/workload/cifar10_cnn.py index 50d377ee..9035af26 100644 --- a/init2winit/projects/optlrschedule/workload/cifar10_cnn.py +++ b/init2winit/projects/optlrschedule/workload/cifar10_cnn.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/projects/optlrschedule/workload/datasets/wikitext_103.py b/init2winit/projects/optlrschedule/workload/datasets/wikitext_103.py index 6363507d..022ea892 100644 --- a/init2winit/projects/optlrschedule/workload/datasets/wikitext_103.py +++ b/init2winit/projects/optlrschedule/workload/datasets/wikitext_103.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/projects/optlrschedule/workload/linear_regression.py b/init2winit/projects/optlrschedule/workload/linear_regression.py index 3fe8b578..8d7a1cee 100644 --- a/init2winit/projects/optlrschedule/workload/linear_regression.py +++ b/init2winit/projects/optlrschedule/workload/linear_regression.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/projects/optlrschedule/workload/optimizers.py b/init2winit/projects/optlrschedule/workload/optimizers.py index 822c5ff9..3afe0840 100644 --- a/init2winit/projects/optlrschedule/workload/optimizers.py +++ b/init2winit/projects/optlrschedule/workload/optimizers.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/projects/optlrschedule/workload/test_workload.py b/init2winit/projects/optlrschedule/workload/test_workload.py index badecb29..4df4c68f 100644 --- a/init2winit/projects/optlrschedule/workload/test_workload.py +++ b/init2winit/projects/optlrschedule/workload/test_workload.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/projects/optlrschedule/workload/wikitext103_transformer.py b/init2winit/projects/optlrschedule/workload/wikitext103_transformer.py index a8f1b8af..c1a9557a 100644 --- a/init2winit/projects/optlrschedule/workload/wikitext103_transformer.py +++ b/init2winit/projects/optlrschedule/workload/wikitext103_transformer.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/projects/optlrschedule/workload/workloads.py b/init2winit/projects/optlrschedule/workload/workloads.py index f81a32c7..85da840d 100644 --- a/init2winit/projects/optlrschedule/workload/workloads.py +++ b/init2winit/projects/optlrschedule/workload/workloads.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/schedules.py b/init2winit/schedules.py index 9b9a80b3..8d27c978 100644 --- a/init2winit/schedules.py +++ b/init2winit/schedules.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/shared_test_utilities.py b/init2winit/shared_test_utilities.py index f88fbda2..bc4ecd7c 100644 --- a/init2winit/shared_test_utilities.py +++ b/init2winit/shared_test_utilities.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/test_checkpoint.py b/init2winit/test_checkpoint.py index cc2d74a4..81a0eb6d 100644 --- a/init2winit/test_checkpoint.py +++ b/init2winit/test_checkpoint.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/test_hyperparameters.py b/init2winit/test_hyperparameters.py index 5ba18064..1a7e5158 100644 --- a/init2winit/test_hyperparameters.py +++ b/init2winit/test_hyperparameters.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/test_schedules.py b/init2winit/test_schedules.py index 68726ea5..6a1574fb 100644 --- a/init2winit/test_schedules.py +++ b/init2winit/test_schedules.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/test_training_metrics_grabber.py b/init2winit/test_training_metrics_grabber.py index cd49f80e..8cfdaf6f 100644 --- a/init2winit/test_training_metrics_grabber.py +++ b/init2winit/test_training_metrics_grabber.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/test_utils.py b/init2winit/test_utils.py index 1b6076a6..f6082c59 100644 --- a/init2winit/test_utils.py +++ b/init2winit/test_utils.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/tools/inspect_dataset.py b/init2winit/tools/inspect_dataset.py index 23c598d4..cbca33af 100644 --- a/init2winit/tools/inspect_dataset.py +++ b/init2winit/tools/inspect_dataset.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/tools/preprocess_ogbg_molpcba.py b/init2winit/tools/preprocess_ogbg_molpcba.py new file mode 100644 index 00000000..3d0db274 --- /dev/null +++ b/init2winit/tools/preprocess_ogbg_molpcba.py @@ -0,0 +1,105 @@ +# coding=utf-8 +# Copyright 2026 The init2winit Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +r"""Tool to preprocess OGBG MolPCBA dataset. + +This script loads the OGBG MolPCBA dataset using the existing pipeline (which +does dynamic batching and padding), unrolls it for a specified number of steps, +and saves the resulting dataset to disk using tf.data.Dataset.save. + +This allows for much faster loading during training by avoiding the expensive +graph processing and dynamic batching at runtime. +""" + +import logging +from absl import app +from absl import flags +from init2winit.dataset_lib import ogbg_molpcba +from init2winit.dataset_lib import ogbg_molpcba_preprocessed +from ml_collections import config_dict +import tqdm + +# Setup logging +logging.basicConfig( + level=logging.INFO, format='%(levelname)s: %(message)s', force=True +) + +# Setup flags +flags.DEFINE_string('output_path', None, 'Path to save the processed dataset.') +flags.DEFINE_integer('num_steps', 10000000, 'Number of steps to generate.') +flags.DEFINE_integer( + 'batch_size', + 256, + 'Batch size (per-host if sharded later, but here it is just the batch size' + ' produced).', +) +flags.DEFINE_integer( + 'target_num_processes', None, 'Target number of processes.' +) +flags.DEFINE_float('batch_nodes_multiplier', 1.0, 'Multiplier for max nodes.') +flags.DEFINE_float('batch_edges_multiplier', 2.0, 'Multiplier for max edges.') +flags.DEFINE_float('avg_nodes_per_graph', 26.0, 'Average nodes per graph.') +flags.DEFINE_float('avg_edges_per_graph', 28.0, 'Average edges per graph.') +flags.DEFINE_bool('add_bidirectional_edges', False, 'Add bidirectional edges.') +flags.DEFINE_bool('add_virtual_node', False, 'Add virtual node.') +flags.DEFINE_bool('add_self_loops', False, 'Add self loops.') +flags.DEFINE_integer('seed', 0, 'Random seed.') + +FLAGS = flags.FLAGS + + +def main(argv): + if len(argv) > 1: + raise app.UsageError('Too many command-line arguments.') + + if not FLAGS.output_path: + raise ValueError('Must specify --output_path') + + # Construct HParams + # We use the defaults from ogbg_molpcba and override with flags + hps = config_dict.ConfigDict(ogbg_molpcba.DEFAULT_HPARAMS) + hps.batch_nodes_multiplier = FLAGS.batch_nodes_multiplier + hps.batch_edges_multiplier = FLAGS.batch_edges_multiplier + hps.avg_nodes_per_graph = int(FLAGS.avg_nodes_per_graph) + hps.avg_edges_per_graph = int(FLAGS.avg_edges_per_graph) + hps.add_bidirectional_edges = FLAGS.add_bidirectional_edges + hps.add_virtual_node = FLAGS.add_virtual_node + hps.add_self_loops = FLAGS.add_self_loops + + logging.info( + 'Generating dataset with batch_size=%d for %d devices for %d steps...', + FLAGS.batch_size, + FLAGS.target_num_processes, + FLAGS.num_steps, + ) + logging.info('Saving to %s', FLAGS.output_path) + + ogbg_molpcba_preprocessed.generate_and_save_dataset( + hps=hps, + output_path=FLAGS.output_path, + num_steps=FLAGS.num_steps, + batch_size=FLAGS.batch_size, + seed=FLAGS.seed, + progress_bar_fn=lambda x: tqdm.tqdm( + x, desc='Generating batches', unit='batch' + ), + target_num_processes=FLAGS.target_num_processes, + ) + + logging.info('Done.') + + +if __name__ == '__main__': + app.run(main) diff --git a/init2winit/trainer_lib/base_trainer.py b/init2winit/trainer_lib/base_trainer.py index c439464b..4df39624 100644 --- a/init2winit/trainer_lib/base_trainer.py +++ b/init2winit/trainer_lib/base_trainer.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit Authors. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -670,7 +670,7 @@ def train(self): # directly in the top-level for loop). batch = next(train_iter) batch = self.finalize_batch_fn(batch) - if self._global_step == 0: + if False and self._global_step == 0: batch_size_pytree = self.get_batch_size(batch) if any( bsz != self._hps.batch_size diff --git a/init2winit/trainer_lib/test_trainer.py b/init2winit/trainer_lib/test_trainer.py index 87e324e8..77020414 100644 --- a/init2winit/trainer_lib/test_trainer.py +++ b/init2winit/trainer_lib/test_trainer.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/trainer_lib/trainer.py b/init2winit/trainer_lib/trainer.py index 187720e1..2befe503 100644 --- a/init2winit/trainer_lib/trainer.py +++ b/init2winit/trainer_lib/trainer.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/trainer_lib/trainer_utils.py b/init2winit/trainer_lib/trainer_utils.py index ec57415d..566b7148 100644 --- a/init2winit/trainer_lib/trainer_utils.py +++ b/init2winit/trainer_lib/trainer_utils.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/trainer_lib/trainers.py b/init2winit/trainer_lib/trainers.py index a5f75608..ae75177a 100644 --- a/init2winit/trainer_lib/trainers.py +++ b/init2winit/trainer_lib/trainers.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/trainer_lib/training_algorithm.py b/init2winit/trainer_lib/training_algorithm.py index e0a8ca7f..a7e473d9 100644 --- a/init2winit/trainer_lib/training_algorithm.py +++ b/init2winit/trainer_lib/training_algorithm.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/trainer_lib/training_algorithms.py b/init2winit/trainer_lib/training_algorithms.py index 2ff2cf67..3c28a59e 100644 --- a/init2winit/trainer_lib/training_algorithms.py +++ b/init2winit/trainer_lib/training_algorithms.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/training_metrics_grabber.py b/init2winit/training_metrics_grabber.py index 38c952ef..8e0335e8 100644 --- a/init2winit/training_metrics_grabber.py +++ b/init2winit/training_metrics_grabber.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit 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/init2winit/utils.py b/init2winit/utils.py index e307e12d..261c9096 100644 --- a/init2winit/utils.py +++ b/init2winit/utils.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The init2winit Authors. +# Copyright 2026 The init2winit Authors. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License.