Remove uses of deprecated jax.lax.tie_in

This has been a no-op since jax v0.2.0, and passes the second argument through unchanged. `tie_in` will be deprecated as of jax v0.4.24; see jax-ml/jax#19413 PiperOrigin-RevId: 602456018
google · Feb 1, 2024 · cfde3a4 · cfde3a4
1 parent d72bd65
commit cfde3a4
Show file tree

Hide file tree

Showing 315 changed files with 316 additions and 316 deletions.
diff --git a/.github/workflows/build.yaml b/.github/workflows/build.yaml
@@ -1,4 +1,4 @@
-# Copyright 2023 The Trax Authors.
+# Copyright 2024 The Trax Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/.readthedocs.yaml b/.readthedocs.yaml
@@ -1,4 +1,4 @@
-# Copyright 2023 The Trax Authors.
+# Copyright 2024 The Trax Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/docs/.readthedocs.yaml b/docs/.readthedocs.yaml
@@ -1,4 +1,4 @@
-# Copyright 2023 The Trax Authors.
+# Copyright 2024 The Trax Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/docs/source/conf.py b/docs/source/conf.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2023 The Trax Authors.
+# Copyright 2024 The Trax Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/oss_scripts/oss_pip_install.sh b/oss_scripts/oss_pip_install.sh
@@ -1,4 +1,4 @@
-# Copyright 2023 The Trax Authors.
+# Copyright 2024 The Trax Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/oss_scripts/oss_release.sh b/oss_scripts/oss_release.sh
@@ -1,4 +1,4 @@
-# Copyright 2023 The Trax Authors.
+# Copyright 2024 The Trax Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/oss_scripts/oss_tests.sh b/oss_scripts/oss_tests.sh
@@ -1,4 +1,4 @@
-# Copyright 2023 The Trax Authors.
+# Copyright 2024 The Trax Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/setup.py b/setup.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2023 The Trax Authors.
+# Copyright 2024 The Trax Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/trax/__init__.py b/trax/__init__.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2023 The Trax Authors.
+# Copyright 2024 The Trax Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/trax/data/__init__.py b/trax/data/__init__.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2023 The Trax Authors.
+# Copyright 2024 The Trax Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/trax/data/debug_data_pipeline.py b/trax/data/debug_data_pipeline.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2023 The Trax Authors.
+# Copyright 2024 The Trax Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/trax/data/inputs.py b/trax/data/inputs.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2023 The Trax Authors.
+# Copyright 2024 The Trax Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/trax/data/inputs_test.py b/trax/data/inputs_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2023 The Trax Authors.
+# Copyright 2024 The Trax Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/trax/data/text_encoder.py b/trax/data/text_encoder.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2023 The Trax Authors.
+# Copyright 2024 The Trax Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/trax/data/text_encoder_build_subword.py b/trax/data/text_encoder_build_subword.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2023 The Trax Authors.
+# Copyright 2024 The Trax Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/trax/data/text_encoder_test.py b/trax/data/text_encoder_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2023 The Trax Authors.
+# Copyright 2024 The Trax Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/trax/data/tf_inputs.py b/trax/data/tf_inputs.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2023 The Trax Authors.
+# Copyright 2024 The Trax Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/trax/data/tf_inputs_test.py b/trax/data/tf_inputs_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2023 The Trax Authors.
+# Copyright 2024 The Trax Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/trax/data/tokenizer.py b/trax/data/tokenizer.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2023 The Trax Authors.
+# Copyright 2024 The Trax Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/trax/data/tokenizer_test.py b/trax/data/tokenizer_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2023 The Trax Authors.
+# Copyright 2024 The Trax Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/trax/fastmath/__init__.py b/trax/fastmath/__init__.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2023 The Trax Authors.
+# Copyright 2024 The Trax Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/trax/fastmath/jax.py b/trax/fastmath/jax.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2023 The Trax Authors.
+# Copyright 2024 The Trax Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/trax/fastmath/numpy.py b/trax/fastmath/numpy.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2023 The Trax Authors.
+# Copyright 2024 The Trax Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/trax/fastmath/ops.py b/trax/fastmath/ops.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2023 The Trax Authors.
+# Copyright 2024 The Trax Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/trax/fastmath/ops_test.py b/trax/fastmath/ops_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2023 The Trax Authors.
+# Copyright 2024 The Trax Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/trax/fastmath/tf.py b/trax/fastmath/tf.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2023 The Trax Authors.
+# Copyright 2024 The Trax Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/trax/import_test.py b/trax/import_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2023 The Trax Authors.
+# Copyright 2024 The Trax Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/trax/jaxboard.py b/trax/jaxboard.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2023 The Trax Authors.
+# Copyright 2024 The Trax Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/trax/layers/__init__.py b/trax/layers/__init__.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2023 The Trax Authors.
+# Copyright 2024 The Trax Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/trax/layers/acceleration.py b/trax/layers/acceleration.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2023 The Trax Authors.
+# Copyright 2024 The Trax Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/trax/layers/acceleration_test.py b/trax/layers/acceleration_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2023 The Trax Authors.
+# Copyright 2024 The Trax Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/trax/layers/activation_fns.py b/trax/layers/activation_fns.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2023 The Trax Authors.
+# Copyright 2024 The Trax Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/trax/layers/activation_fns_test.py b/trax/layers/activation_fns_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2023 The Trax Authors.
+# Copyright 2024 The Trax Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/trax/layers/assert_shape.py b/trax/layers/assert_shape.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2023 The Trax Authors.
+# Copyright 2024 The Trax Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/trax/layers/assert_shape_test.py b/trax/layers/assert_shape_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2023 The Trax Authors.
+# Copyright 2024 The Trax Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/trax/layers/attention.py b/trax/layers/attention.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2023 The Trax Authors.
+# Copyright 2024 The Trax Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/trax/layers/attention_test.py b/trax/layers/attention_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2023 The Trax Authors.
+# Copyright 2024 The Trax Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/trax/layers/base.py b/trax/layers/base.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2023 The Trax Authors.
+# Copyright 2024 The Trax Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/trax/layers/base_test.py b/trax/layers/base_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2023 The Trax Authors.
+# Copyright 2024 The Trax Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/trax/layers/combinators.py b/trax/layers/combinators.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2023 The Trax Authors.
+# Copyright 2024 The Trax Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/trax/layers/combinators_test.py b/trax/layers/combinators_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2023 The Trax Authors.
+# Copyright 2024 The Trax Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/trax/layers/convolution.py b/trax/layers/convolution.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2023 The Trax Authors.
+# Copyright 2024 The Trax Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/trax/layers/convolution_test.py b/trax/layers/convolution_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2023 The Trax Authors.
+# Copyright 2024 The Trax Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/trax/layers/core.py b/trax/layers/core.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2023 The Trax Authors.
+# Copyright 2024 The Trax Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/trax/layers/core_test.py b/trax/layers/core_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2023 The Trax Authors.
+# Copyright 2024 The Trax Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/trax/layers/deconvolution.py b/trax/layers/deconvolution.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2023 The Trax Authors.
+# Copyright 2024 The Trax Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/trax/layers/deconvolution_test.py b/trax/layers/deconvolution_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2023 The Trax Authors.
+# Copyright 2024 The Trax Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/trax/layers/initializers.py b/trax/layers/initializers.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2023 The Trax Authors.
+# Copyright 2024 The Trax Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/trax/layers/initializers_test.py b/trax/layers/initializers_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2023 The Trax Authors.
+# Copyright 2024 The Trax Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/trax/layers/metrics.py b/trax/layers/metrics.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2023 The Trax Authors.
+# Copyright 2024 The Trax Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/trax/layers/metrics_test.py b/trax/layers/metrics_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2023 The Trax Authors.
+# Copyright 2024 The Trax Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/trax/layers/normalization.py b/trax/layers/normalization.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2023 The Trax Authors.
+# Copyright 2024 The Trax Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/trax/layers/normalization_test.py b/trax/layers/normalization_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2023 The Trax Authors.
+# Copyright 2024 The Trax Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/trax/layers/pooling.py b/trax/layers/pooling.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2023 The Trax Authors.
+# Copyright 2024 The Trax Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/trax/layers/pooling_test.py b/trax/layers/pooling_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2023 The Trax Authors.
+# Copyright 2024 The Trax Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/trax/layers/research/__init__.py b/trax/layers/research/__init__.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2023 The Trax Authors.
+# Copyright 2024 The Trax Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/trax/layers/research/efficient_attention.py b/trax/layers/research/efficient_attention.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2023 The Trax Authors.
+# Copyright 2024 The Trax Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/trax/layers/research/efficient_attention_test.py b/trax/layers/research/efficient_attention_test.py
@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2023 The Trax Authors.
+# Copyright 2024 The Trax Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.