tf-keras-nightly 2.20.0.dev2025062209__py3-none-any.whl → 2.20.0.dev2025082909__py3-none-any.whl

tf_keras/src/layers/preprocessing/benchmarks/image_preproc_benchmark.py

@@ -1,159 +0,0 @@
-# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
-#
-# 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.
-# ==============================================================================
-"""Benchmark for TF-Keras image preprocessing layer."""
-
-import functools
-import time
-
-import numpy as np
-import tensorflow.compat.v2 as tf
-
-import tf_keras.src as keras
-from tf_keras.src.layers.preprocessing import image_preprocessing
-
-LOWER = 0.2
-UPPER = 0.4
-BATCH_SIZE = 32
-
-
-def rotate(inputs):
-    """rotate image."""
-    inputs_shape = tf.shape(inputs)
-    batch_size = inputs_shape[0]
-    img_hd = tf.cast(inputs_shape[1], tf.float32)
-    img_wd = tf.cast(inputs_shape[2], tf.float32)
-    min_angle = LOWER * 2.0 * np.pi
-    max_angle = UPPER * 2.0 * np.pi
-    angles = tf.random.uniform(
-        shape=[batch_size], minval=min_angle, maxval=max_angle
-    )
-    return image_preprocessing.transform(
-        inputs, image_preprocessing.get_rotation_matrix(angles, img_hd, img_wd)
-    )
-
-
-def zoom(inputs):
-    """zoom image."""
-    inputs_shape = tf.shape(inputs)
-    batch_size = inputs_shape[0]
-    img_hd = tf.cast(inputs_shape[1], tf.float32)
-    img_wd = tf.cast(inputs_shape[2], tf.float32)
-    height_zoom = tf.random.uniform(
-        shape=[batch_size, 1], minval=1.0 + LOWER, maxval=1.0 + UPPER
-    )
-    width_zoom = tf.random.uniform(
-        shape=[batch_size, 1], minval=1.0 + LOWER, maxval=1.0 + UPPER
-    )
-    zooms = tf.cast(
-        tf.concat([width_zoom, height_zoom], axis=1), dtype=tf.float32
-    )
-    return image_preprocessing.transform(
-        inputs, image_preprocessing.get_zoom_matrix(zooms, img_hd, img_wd)
-    )
-
-
-def image_augmentation(inputs, batch_size):
-    """image augmentation."""
-    img = inputs
-    img = tf.image.resize(img, size=[224, 224])
-    img = tf.image.random_crop(img, size=[batch_size, 224, 224, 3])
-    img = rotate(img)
-    img = zoom(img)
-    return img
-
-
-class BenchmarkLayer(tf.test.Benchmark):
-    """Benchmark the layer forward pass."""
-
-    def run_dataset_implementation(self, batch_size):
-        num_repeats = 5
-        starts = []
-        ends = []
-        for _ in range(num_repeats):
-            ds = tf.data.Dataset.from_tensor_slices(
-                np.random.random((batch_size, 256, 256, 3))
-            )
-            ds = ds.shuffle(batch_size * 100)
-            ds = ds.batch(batch_size)
-            ds = ds.prefetch(batch_size)
-            img_augmentation = functools.partial(
-                image_augmentation, batch_size=batch_size
-            )
-            ds = ds.map(img_augmentation, num_parallel_calls=8)
-            starts.append(time.time())
-            count = 0
-            # Benchmarked code begins here.
-            for i in ds:
-                _ = i
-                count += 1
-            # Benchmarked code ends here.
-            ends.append(time.time())
-
-        avg_time = np.mean(np.array(ends) - np.array(starts)) / count
-        return avg_time
-
-    def bm_layer_implementation(self, batch_size):
-        with tf.device("/gpu:0"):
-            img = keras.Input(shape=(256, 256, 3), dtype=tf.float32)
-            preprocessor = keras.Sequential(
-                [
-                    image_preprocessing.Resizing(224, 224),
-                    image_preprocessing.RandomCrop(height=224, width=224),
-                    image_preprocessing.RandomRotation(factor=(0.2, 0.4)),
-                    image_preprocessing.RandomFlip(mode="horizontal"),
-                    image_preprocessing.RandomZoom(0.2, 0.2),
-                ]
-            )
-            _ = preprocessor(img)
-
-            num_repeats = 5
-            starts = []
-            ends = []
-            for _ in range(num_repeats):
-                ds = tf.data.Dataset.from_tensor_slices(
-                    np.random.random((batch_size, 256, 256, 3))
-                )
-                ds = ds.shuffle(batch_size * 100)
-                ds = ds.batch(batch_size)
-                ds = ds.prefetch(batch_size)
-                starts.append(time.time())
-                count = 0
-                # Benchmarked code begins here.
-                for i in ds:
-                    _ = preprocessor(i)
-                    count += 1
-                # Benchmarked code ends here.
-                ends.append(time.time())
-
-        avg_time = np.mean(np.array(ends) - np.array(starts)) / count
-        name = f"image_preprocessing|batch_{batch_size}"
-        baseline = self.run_dataset_implementation(batch_size)
-        extras = {
-            "dataset implementation baseline": baseline,
-            "delta seconds": (baseline - avg_time),
-            "delta percent": ((baseline - avg_time) / baseline) * 100,
-        }
-        self.report_benchmark(
-            iters=num_repeats, wall_time=avg_time, extras=extras, name=name
-        )
-
-    def benchmark_vocab_size_by_batch(self):
-        for batch in [32, 64, 256]:
-            self.bm_layer_implementation(batch_size=batch)
-
-
-if __name__ == "__main__":
-    tf.test.main()
-

tf_keras/src/layers/preprocessing/benchmarks/index_lookup_adapt_benchmark.py

@@ -1,135 +0,0 @@
-# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
-#
-# 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.
-# ==============================================================================
-"""Benchmark for TF-Keras text vectorization preprocessing layer's adapt method.
-"""
-
-import collections
-import itertools
-import random
-import string
-import time
-
-import numpy as np
-import tensorflow.compat.v2 as tf
-
-import tf_keras.src as keras
-from tf_keras.src.layers.preprocessing import index_lookup
-
-tf.compat.v1.enable_v2_behavior()
-
-
-# word_gen creates random sequences of ASCII letters (both lowercase and upper).
-# The number of unique strings is ~2,700.
-def word_gen():
-    for _ in itertools.count(1):
-        yield "".join(random.choice(string.ascii_letters) for i in range(2))
-
-
-def get_top_k(dataset, k):
-    """Python implementation of vocabulary building using a defaultdict."""
-    counts = collections.defaultdict(int)
-    for tensor in dataset:
-        data = tensor.numpy()
-        for element in data:
-            counts[element] += 1
-    sorted_vocab = [
-        k
-        for k, _ in sorted(
-            counts.items(), key=lambda item: item[1], reverse=True
-        )
-    ]
-    if len(sorted_vocab) > k:
-        sorted_vocab = sorted_vocab[:k]
-    return sorted_vocab
-
-
-class BenchmarkAdapt(tf.test.Benchmark):
-    """Benchmark adapt."""
-
-    def run_numpy_implementation(self, num_elements, batch_size, k):
-        """Test the python implementation."""
-        ds = tf.data.Dataset.from_generator(
-            word_gen, tf.string, tf.TensorShape([])
-        )
-        batched_ds = ds.take(num_elements).batch(batch_size)
-        input_t = keras.Input(shape=(), dtype=tf.string)
-        layer = index_lookup.IndexLookup(
-            max_tokens=k,
-            num_oov_indices=0,
-            mask_token=None,
-            oov_token="OOV",
-            dtype=tf.string,
-        )
-        _ = layer(input_t)
-        num_repeats = 5
-        starts = []
-        ends = []
-        for _ in range(num_repeats):
-            starts.append(time.time())
-            vocab = get_top_k(batched_ds, k)
-            layer.set_vocabulary(vocab)
-            ends.append(time.time())
-        avg_time = np.mean(np.array(ends) - np.array(starts))
-        return avg_time
-
-    def bm_adapt_implementation(self, num_elements, batch_size, k):
-        """Test the KPL adapt implementation."""
-        ds = tf.data.Dataset.from_generator(
-            word_gen, tf.string, tf.TensorShape([])
-        )
-        batched_ds = ds.take(num_elements).batch(batch_size)
-        input_t = keras.Input(shape=(), dtype=tf.string)
-        layer = index_lookup.IndexLookup(
-            max_tokens=k,
-            num_oov_indices=0,
-            mask_token=None,
-            oov_token="OOV",
-            dtype=tf.string,
-        )
-        _ = layer(input_t)
-        num_repeats = 5
-        starts = []
-        ends = []
-        for _ in range(num_repeats):
-            starts.append(time.time())
-            layer.adapt(batched_ds)
-            ends.append(time.time())
-        avg_time = np.mean(np.array(ends) - np.array(starts))
-        name = "index_lookup_adapt|%s_elements|vocab_size_%s|batch_%s" % (
-            num_elements,
-            k,
-            batch_size,
-        )
-        baseline = self.run_numpy_implementation(num_elements, batch_size, k)
-        extras = {
-            "numpy implementation baseline": baseline,
-            "delta seconds": (baseline - avg_time),
-            "delta percent": ((baseline - avg_time) / baseline) * 100,
-        }
-        self.report_benchmark(
-            iters=num_repeats, wall_time=avg_time, extras=extras, name=name
-        )
-
-    def benchmark_vocab_size_by_batch(self):
-        for vocab_size in [100, 1000, 10000, 100000, 1000000]:
-            for batch in [1, 16, 2048]:
-                self.bm_adapt_implementation(
-                    vocab_size, batch, int(vocab_size / 10)
-                )
-
-
-if __name__ == "__main__":
-    tf.test.main()
-

tf_keras/src/layers/preprocessing/benchmarks/index_lookup_forward_benchmark.py

@@ -1,144 +0,0 @@
-# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
-#
-# 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.
-# ==============================================================================
-"""Benchmark for TF-Keras text vectorization preprocessing layer's adapt method.
-"""
-
-import os
-import random
-import string
-import time
-
-import numpy as np
-import tensorflow.compat.v2 as tf
-
-import tf_keras.src as keras
-from tf_keras.src.layers.preprocessing import index_lookup
-
-
-# word_gen creates random sequences of ASCII letters (both lowercase and upper).
-# The number of unique strings is ~2,700.
-def tensor_gen(batch, num_elements):
-    data = []
-    for _ in range(batch):
-        batch_element = []
-        for _ in range(num_elements - 1):
-            tok = "".join(random.choice(string.ascii_letters) for i in range(2))
-            batch_element.append(tok)
-        batch_element.append("")  # Explicitly test the empty string.
-        data.append(batch_element)
-    return tf.constant(data)
-
-
-def get_vocab():
-    vocab = list(
-        set([a + b for a in string.ascii_letters for b in string.ascii_letters])
-    )
-    vocab.sort()
-    return vocab
-
-
-# This class uses TestCase for get_temp_dir().
-class BenchmarkLookup(tf.test.Benchmark):
-    """Benchmark the index lookup layer's forward pass."""
-
-    def _write_to_temp_file(self, file_name, vocab_list):
-        vocab_path = os.path.join(self.get_temp_dir(), file_name + ".txt")
-        with tf.io.gfile.GFile(vocab_path, "w") as writer:
-            for vocab in vocab_list:
-                writer.write(vocab + "\n")
-            writer.flush()
-            writer.close()
-        return vocab_path
-
-    def run_numpy_implementation(self, data, vocab):
-        """Test the python implementation."""
-        input_t = keras.Input(shape=(), dtype=tf.string)
-        layer = index_lookup.IndexLookup(
-            vocabulary=vocab,
-            max_tokens=None,
-            num_oov_indices=1,
-            mask_token="",
-            oov_token="OOV",
-            dtype=tf.string,
-        )
-        out_t = layer(input_t)
-        model = keras.Model(input_t, out_t)
-        num_repeats = 5
-        starts = []
-        ends = []
-        _ = model(data)
-        for _ in range(num_repeats):
-            starts.append(time.time())
-            out = model(data)
-            ends.append(time.time())
-        avg_time = np.mean(np.array(ends) - np.array(starts))
-        return avg_time, out
-
-    def bm_adapt_implementation(self, num_elements, batch_size):
-        """Test the KPL adapt implementation."""
-        vocab = get_vocab()
-        vocab_file = self._write_to_temp_file("vocab", vocab)
-        vocabulary_initializer = tf.lookup.TextFileInitializer(
-            filename=vocab_file,
-            key_dtype=tf.string,
-            key_index=tf.lookup.TextFileIndex.WHOLE_LINE,
-            value_dtype=tf.int64,
-            value_index=tf.lookup.TextFileIndex.LINE_NUMBER,
-            value_index_offset=2,
-        )
-        input_t = keras.Input(shape=(), dtype=tf.string)
-        layer = index_lookup.IndexLookup(
-            vocabulary=vocabulary_initializer,
-            max_tokens=None,
-            num_oov_indices=1,
-            mask_token="",
-            oov_token="OOV",
-            dtype=tf.string,
-        )
-        out_t = layer(input_t)
-        model = keras.Model(input_t, out_t)
-        num_repeats = 5
-        starts = []
-        ends = []
-        data = tensor_gen(batch_size, num_elements)
-        _ = model(data)
-        for _ in range(num_repeats):
-            starts.append(time.time())
-            _ = model(data)
-            ends.append(time.time())
-        avg_time = np.mean(np.array(ends) - np.array(starts))
-        baseline, _ = self.run_numpy_implementation(data, vocab)
-        extras = {
-            "numpy implementation baseline": baseline,
-            "delta seconds": (baseline - avg_time),
-            "delta percent": ((baseline - avg_time) / baseline) * 100,
-        }
-        name = "index_lookup_forward|%s_elements|batch_%s" % (
-            num_elements,
-            batch_size,
-        )
-        self.report_benchmark(
-            iters=num_repeats, wall_time=avg_time, extras=extras, name=name
-        )
-
-    def benchmark_vocab_size_by_batch(self):
-        for tensor_size in [100, 1000, 10000]:
-            for batch in [1, 16, 2048]:
-                self.bm_adapt_implementation(tensor_size, batch)
-
-
-if __name__ == "__main__":
-    tf.test.main()
-

tf_keras/src/layers/preprocessing/benchmarks/normalization_adapt_benchmark.py

@@ -1,124 +0,0 @@
-# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
-#
-# 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.
-# ==============================================================================
-"""Benchmark for TF-Keras text vectorization preprocessing layer's adapt method.
-"""
-
-import time
-
-import numpy as np
-import tensorflow.compat.v2 as tf
-
-import tf_keras.src as keras
-from tf_keras.src.layers.preprocessing import normalization
-
-
-def reduce_fn(state, values):
-    """tf.data.Dataset-friendly implementation of mean and variance."""
-    k, n, ex, ex2 = state
-    # If this is the first iteration, we pick the first value to be 'k',
-    # which helps with precision - we assume that k is close to an average
-    # value and calculate mean and variance with respect to that.
-    k = tf.cond(tf.equal(n, 0), lambda: values[0], lambda: k)
-
-    sum_v = tf.reduce_sum(values, axis=0)
-    sum_v2 = tf.reduce_sum(tf.square(values), axis=0)
-    ones = tf.ones_like(values, dtype=tf.int32)
-    batch_size = tf.reduce_sum(ones, axis=0)
-    batch_size_f = tf.cast(batch_size, tf.float32)
-
-    ex = 0 + sum_v - tf.multiply(batch_size_f, k)
-    ex2 = (
-        0
-        + sum_v2
-        + tf.multiply(
-            batch_size_f,
-            (tf.square(k) - tf.multiply(tf.multiply(2.0, k), sum_v)),
-        )
-    )
-
-    return (k, n + batch_size, ex, ex2)
-
-
-class BenchmarkAdapt(tf.test.Benchmark):
-    """Benchmark adapt."""
-
-    def run_dataset_implementation(self, num_elements, batch_size):
-        input_t = keras.Input(shape=(1,))
-        layer = normalization.Normalization()
-        _ = layer(input_t)
-
-        num_repeats = 5
-        starts = []
-        ends = []
-        for _ in range(num_repeats):
-            ds = tf.data.Dataset.range(num_elements)
-            ds = ds.map(lambda x: tf.expand_dims(tf.cast(x, tf.float32), -1))
-            ds = ds.batch(batch_size)
-
-            starts.append(time.time())
-            # Benchmarked code begins here.
-            k, n, ex, ex2 = ds.reduce((0.0, 0, 0.0, 0.0), reduce_fn)
-            mean = k.numpy() + ex.numpy() / n.numpy()
-            var = (ex2.numpy() - (ex.numpy() * ex.numpy()) / n.numpy()) / (
-                n.numpy() - 1
-            )
-            layer.set_weights([mean, var])
-            # Benchmarked code ends here.
-            ends.append(time.time())
-
-        avg_time = np.mean(np.array(ends) - np.array(starts))
-        return avg_time
-
-    def bm_adapt_implementation(self, num_elements, batch_size):
-        """Test the KPL adapt implementation."""
-        input_t = keras.Input(shape=(1,), dtype=tf.float32)
-        layer = normalization.Normalization()
-        _ = layer(input_t)
-
-        num_repeats = 5
-        starts = []
-        ends = []
-        for _ in range(num_repeats):
-            ds = tf.data.Dataset.range(num_elements)
-            ds = ds.map(lambda x: tf.expand_dims(tf.cast(x, tf.float32), -1))
-            ds = ds.batch(batch_size)
-
-            starts.append(time.time())
-            # Benchmarked code begins here.
-            layer.adapt(ds)
-            # Benchmarked code ends here.
-            ends.append(time.time())
-
-        avg_time = np.mean(np.array(ends) - np.array(starts))
-        name = f"normalization_adapt|{num_elements}_elements|batch_{batch_size}"
-        baseline = self.run_dataset_implementation(num_elements, batch_size)
-        extras = {
-            "tf.data implementation baseline": baseline,
-            "delta seconds": (baseline - avg_time),
-            "delta percent": ((baseline - avg_time) / baseline) * 100,
-        }
-        self.report_benchmark(
-            iters=num_repeats, wall_time=avg_time, extras=extras, name=name
-        )
-
-    def benchmark_vocab_size_by_batch(self):
-        for vocab_size in [100, 1000, 10000, 100000, 1000000]:
-            for batch in [1, 16, 2048]:
-                self.bm_adapt_implementation(vocab_size, batch)
-
-
-if __name__ == "__main__":
-    tf.test.main()
-

tf_keras/src/layers/preprocessing/benchmarks/weighted_embedding_varlen_benchmark.py

@@ -1,99 +0,0 @@
-# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
-#
-# 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.
-# ==============================================================================
-"""Benchmark for KPL implementation of weighted embedding column with
-varying-length inputs."""
-
-import tensorflow.compat.v2 as tf
-
-import tf_keras.src as keras
-from tf_keras.src.layers.preprocessing.benchmarks import (
-    feature_column_benchmark as fc_bm,
-)
-
-# isort: off
-from tensorflow.python.eager.def_function import (
-    function as tf_function,
-)
-
-NUM_REPEATS = 10
-BATCH_SIZES = [32, 256]
-
-
-### KPL AND FC IMPLEMENTATION BENCHMARKS ###
-def embedding_varlen(batch_size, max_length):
-    """Benchmark a variable-length embedding."""
-    # Data and constants.
-    embedding_size = 32768
-    data = fc_bm.create_data(
-        max_length, batch_size * NUM_REPEATS, embedding_size - 1, dtype=int
-    )
-    weight = tf.ones_like(data, dtype=tf.float32)
-
-    # TF-Keras implementation
-    data_input = keras.Input(
-        shape=(None,), ragged=True, name="data", dtype=tf.int64
-    )
-    weight_input = keras.Input(
-        shape=(None,), ragged=True, name="weight", dtype=tf.float32
-    )
-    embedded_data = keras.layers.Embedding(embedding_size, 256)(data_input)
-    weighted_embedding = tf.multiply(
-        embedded_data, tf.expand_dims(weight_input, -1)
-    )
-    reduced_embedding = tf.reduce_sum(weighted_embedding, axis=1)
-    model = keras.Model([data_input, weight_input], reduced_embedding)
-
-    # FC implementation
-    fc = tf.feature_column.embedding_column(
-        tf.feature_column.weighted_categorical_column(
-            tf.feature_column.categorical_column_with_identity(
-                "data", num_buckets=embedding_size - 1
-            ),
-            weight_feature_key="weight",
-        ),
-        dimension=256,
-    )
-
-    # Wrap the FC implementation in a tf.function for a fair comparison
-    @tf_function()
-    def fc_fn(tensors):
-        fc.transform_feature(
-            tf.__internal__.feature_column.FeatureTransformationCache(tensors),
-            None,
-        )
-
-    # Benchmark runs
-    keras_data = {"data": data, "weight": weight}
-    k_avg_time = fc_bm.run_keras(keras_data, model, batch_size, NUM_REPEATS)
-
-    fc_data = {"data": data.to_sparse(), "weight": weight.to_sparse()}
-    fc_avg_time = fc_bm.run_fc(fc_data, fc_fn, batch_size, NUM_REPEATS)
-
-    return k_avg_time, fc_avg_time
-
-
-class BenchmarkLayer(fc_bm.LayerBenchmark):
-    """Benchmark the layer forward pass."""
-
-    def benchmark_layer(self):
-        for batch in BATCH_SIZES:
-            name = f"weighted_embedding|varlen|batch_{batch}"
-            k_time, f_time = embedding_varlen(batch_size=batch, max_length=256)
-            self.report(name, k_time, f_time, NUM_REPEATS)
-
-
-if __name__ == "__main__":
-    tf.test.main()
-

tf_keras/src/saving/legacy/saved_model/create_test_saved_model.py

@@ -1,37 +0,0 @@
-"""A binary that creates a serialized SavedModel from a keras model.
-
-This is used in tests to ensure that model serialization is deterministic across
-different processes.
-"""
-
-import tensorflow.compat.v2 as tf
-from absl import app
-from absl import flags
-
-from tf_keras.src import regularizers
-from tf_keras.src.testing_infra import test_utils
-
-flags.DEFINE_string("output_path", "", "The path to write the SavedModel at.")
-
-FLAGS = flags.FLAGS
-
-
-def main(_) -> None:
-    with test_utils.model_type_scope("functional"):
-        model = test_utils.get_small_mlp(1, 4, input_dim=3)
-        model.layers[-1].activity_regularizer = regularizers.get("l2")
-        model.activity_regularizer = regularizers.get("l2")
-        model.compile(loss="mse", optimizer="rmsprop")
-
-        def callable_loss():
-            return tf.reduce_sum(model.weights[0])
-
-        model.add_loss(callable_loss)
-
-        print(f"_____Writing saved model to: {FLAGS.output_path}")
-        model.save(FLAGS.output_path)
-
-
-if __name__ == "__main__":
-    app.run(main)
-