jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def test_update_sync_on_read_variable(self, strategy):
    # Updating a synchronizaiton=ON_READ in the replica context should just
    # update the local value. Updating it in the cross replica context updates
    # each component of the variable. Both are true with a loaded model.
    #
    # Note that if assigning a variable whose aggregation=SUM in the cross
    # replica context, each replica is assigned with the value divided by the
    # number of replicas.

    class Model(tf.Module):

      def __init__(self):
        self.v = tf.Variable(
            0.,
            synchronization=tf.VariableSynchronization.ON_READ,
            aggregation=tf.VariableAggregation.SUM)

      @tf.function(input_signature=[tf.TensorSpec(shape=(), dtype=tf.float32)])
      def update(self, value):
        self.v.assign_add(value)

    export_dir = self.get_temp_dir()
    value = strategy.experimental_distribute_values_from_function(
        lambda ctx: tf.identity([3., 7.][ctx.replica_id_in_sync_group]))
    with strategy.scope():
      m = Model()
      tf.saved_model.save(m, export_dir)
      self.evaluate(m.v.assign(10.))
      self.assertAllEqual(
          self.evaluate(strategy.experimental_local_results(m.v)), [5., 5.])
      del m
      # TODO(b/161488560): strategy.run doesn't work with tf.function with
      # input_signature.
      # self.evaluate(strategy.run(m.update, args=(value,)))
      # self.assertAllEqual(
      #     self.evaluate(strategy.experimental_local_results(m.v)), [8., 12.])

    with strategy.scope():
      loaded = tf.saved_model.load(export_dir)
      self.evaluate(loaded.v.assign(10.))
      self.assertAllEqual(
          self.evaluate(strategy.experimental_local_results(loaded.v)),
          [5., 5.])
      self.evaluate(strategy.run(loaded.update, args=(value,)))
      self.assertAllEqual(
          self.evaluate(strategy.experimental_local_results(loaded.v)),
          [8., 12.])
//  [python]
def __init__(self, saved_model_dir_or_frozen_graph: Text, *args, **kwargs):
    """ Initialize the SavedModelDriver.

    Args:
      saved_model_dir_or_frozen_graph: Saved model or frozen graph path.
    """
    # Load saved model if it is a folder.
    super().__init__(*args, **kwargs)
    if tf.saved_model.contains_saved_model(saved_model_dir_or_frozen_graph):
      self.model = tf.compat.v1.saved_model.load(saved_model_dir_or_frozen_graph)
      return

    # Load a frozen graph.
    def wrap_frozen_graph(graph_def, inputs, outputs):
      # https://www.tensorflow.org/guide/migrate
      imports_graph_def_fn = lambda: tf.import_graph_def(graph_def, name='')
      wrapped_import = tf.compat.v1.wrap_function(imports_graph_def_fn, [])
      import_graph = wrapped_import.graph
      return wrapped_import.prune(
          tf.nest.map_structure(import_graph.as_graph_element, inputs),
          tf.nest.map_structure(import_graph.as_graph_element, outputs))

    graph_def = tf.Graph().as_graph_def()
    with tf.io.gfile.GFile(saved_model_dir_or_frozen_graph, 'rb') as f:
      graph_def.ParseFromString(f.read())

    self.model = wrap_frozen_graph(
        graph_def,
        inputs='images:0',
        outputs=['Identity:0', 'Identity_1:0', 'Identity_2:0', 'Identity_3:0'])
//  [python]
def __init__(self, saved_model_dir_or_frozen_graph: Text, *args, **kwargs):
    """ Initialize the SavedModelDriver.

    Args:
      saved_model_dir_or_frozen_graph: Saved model or frozen graph path.
    """
    # Load saved model if it is a folder.
    super().__init__(*args, **kwargs)
    if tf.saved_model.contains_saved_model(saved_model_dir_or_frozen_graph):
      self.model = tf.saved_model.load(saved_model_dir_or_frozen_graph)
      return

    # Load a frozen graph.
    def wrap_frozen_graph(graph_def, inputs, outputs):
      # https://www.tensorflow.org/guide/migrate
      imports_graph_def_fn = lambda: tf.import_graph_def(graph_def, name='')
      wrapped_import = tf.compat.v1.wrap_function(imports_graph_def_fn, [])
      import_graph = wrapped_import.graph
      return wrapped_import.prune(
          tf.nest.map_structure(import_graph.as_graph_element, inputs),
          tf.nest.map_structure(import_graph.as_graph_element, outputs))

    graph_def = tf.Graph().as_graph_def()
    with tf.io.gfile.GFile(saved_model_dir_or_frozen_graph, 'rb') as f:
      graph_def.ParseFromString(f.read())

    self.model = wrap_frozen_graph(
        graph_def,
        inputs='images:0',
        outputs=['Identity:0', 'Identity_1:0', 'Identity_2:0', 'Identity_3:0'])
//  [python]
def test_update_mirrored_variable(self, strategy):
    # This is also uncommon since most model parameters should be updated by
    # optimizer, and this test case is for completeness.
    #
    # In the cross replica context, assigning to the variable assigns the same
    # value to all replicas. This is true with the loaded model as well.
    #
    # However in replica context, MirroredVariable (synchronization=ON_WRITE)
    # in a loaded model behaves differently. Updating MirroredVariable only
    # update the current replica's variable with the current replica's value.
    # There's no aggregation. This doesn't affect variables that are updated
    # through optimizer. This is work as intended but can be surprising.

    class Model(tf.Module):

      def __init__(self):
        self.v = tf.Variable(
            0.,
            synchronization=tf.VariableSynchronization.ON_WRITE,
            aggregation=tf.VariableAggregation.MEAN)

      @tf.function(input_signature=[tf.TensorSpec(shape=(), dtype=tf.float32)])
      def update(self, value):
        self.v.assign_add(value)

    export_dir = self.get_temp_dir()
    value = strategy.experimental_distribute_values_from_function(
        lambda ctx: tf.identity([1., 2.][ctx.replica_id_in_sync_group]))
    with strategy.scope():
      m = Model()
      tf.saved_model.save(m, export_dir)
      del m

    with strategy.scope():
      loaded = tf.compat.v1.saved_model.load(export_dir)
    self.assertAllEqual(
        self.evaluate(strategy.experimental_local_results(loaded.v)), [0., 0.])
    self.evaluate(loaded.v.assign(1.))
    self.assertAllEqual(
        self.evaluate(strategy.experimental_local_results(loaded.v)), [1., 1.])
    strategy.run(loaded.update, args=(value,))
    self.assertAllEqual(
        self.evaluate(strategy.experimental_local_results(loaded.v)), [2., 3.])
//  [python]
def test_update_mirrored_variable(self, strategy):
    # This is also uncommon since most model parameters should be updated by
    # optimizer, and this test case is for completeness.
    #
    # In the cross replica context, assigning to the variable assigns the same
    # value to all replicas. This is true with the loaded model as well.
    #
    # However in replica context, MirroredVariable (synchronization=ON_WRITE)
    # in a loaded model behaves differently. Updating MirroredVariable only
    # update the current replica's variable with the current replica's value.
    # There's no aggregation. This doesn't affect variables that are updated
    # through optimizer. This is work as intended but can be surprising.

    class Model(tf.Module):

      def __init__(self):
        self.v = tf.Variable(
            0.,
            synchronization=tf.VariableSynchronization.ON_WRITE,
            aggregation=tf.VariableAggregation.MEAN)

      @tf.function(input_signature=[tf.TensorSpec(shape=(), dtype=tf.float32)])
      def update(self, value):
        self.v.assign_add(value)

    export_dir = self.get_temp_dir()
    value = strategy.experimental_distribute_values_from_function(
        lambda ctx: tf.identity([1., 2.][ctx.replica_id_in_sync_group]))
    with strategy.scope():
      m = Model()
      tf.saved_model.save(m, export_dir)
      del m

    with strategy.scope():
      loaded = tf.saved_model.load(export_dir)
    self.assertAllEqual(
        self.evaluate(strategy.experimental_local_results(loaded.v)), [0., 0.])
    self.evaluate(loaded.v.assign(1.))
    self.assertAllEqual(
        self.evaluate(strategy.experimental_local_results(loaded.v)), [1., 1.])
    strategy.run(loaded.update, args=(value,))
    self.assertAllEqual(
        self.evaluate(strategy.experimental_local_results(loaded.v)), [2., 3.])
//  [python]
def test_low_level_model_export(self, model_type):
        temp_filepath = os.path.join(self.get_temp_dir(), "exported_model")

        model = get_model(model_type)
        ref_input = tf.random.normal((3, 10))
        ref_output = model(ref_input)

        # Test variable tracking
        export_archive = export_lib.ExportArchive()
        export_archive.track(model)
        self.assertLen(export_archive.variables, 8)
        self.assertLen(export_archive.trainable_variables, 6)
        self.assertLen(export_archive.non_trainable_variables, 2)

        export_archive = export_lib.ExportArchive()
        export_archive.track(model)
        export_archive.add_endpoint(
            "call",
            model.__call__,
            input_signature=[tf.TensorSpec(shape=(None, 10), dtype=tf.float32)],
        )
        export_archive.write_out(temp_filepath)
        revived_model = tf.compat.v1.saved_model.load(temp_filepath)
        self.assertAllClose(ref_output, revived_model.call(ref_input))
        # Test with a different batch size
        revived_model.call(tf.random.normal((6, 10)))
//  [python]
def test_low_level_model_export(self, model_type):
        temp_filepath = os.path.join(self.get_temp_dir(), "exported_model")

        model = get_model(model_type)
        ref_input = tf.random.normal((3, 10))
        ref_output = model(ref_input)

        # Test variable tracking
        export_archive = export_lib.ExportArchive()
        export_archive.track(model)
        self.assertLen(export_archive.variables, 8)
        self.assertLen(export_archive.trainable_variables, 6)
        self.assertLen(export_archive.non_trainable_variables, 2)

        export_archive = export_lib.ExportArchive()
        export_archive.track(model)
        export_archive.add_endpoint(
            "call",
            model.__call__,
            input_signature=[tf.TensorSpec(shape=(None, 10), dtype=tf.float32)],
        )
        export_archive.write_out(temp_filepath)
        revived_model = tf.saved_model.load(temp_filepath)
        self.assertAllClose(ref_output, revived_model.call(ref_input))
        # Test with a different batch size
        revived_model.call(tf.random.normal((6, 10)))
//  [python]
def test_low_level_model_export_with_alias(self):
        temp_filepath = os.path.join(self.get_temp_dir(), "exported_model")

        model = get_model()
        ref_input = tf.random.normal((3, 10))
        ref_output = model(ref_input)

        export_archive = export_lib.ExportArchive()
        export_archive.track(model)
        fn = export_archive.add_endpoint(
            "call",
            model.__call__,
            input_signature=[tf.TensorSpec(shape=(None, 10), dtype=tf.float32)],
        )
        export_archive.write_out(
            temp_filepath,
            tf.saved_model.SaveOptions(function_aliases={"call_alias": fn}),
        )
        revived_model = tf.compat.v1.saved_model.load(
            temp_filepath,
            options=tf.compat.v1.saved_model.LoadOptions(
                experimental_load_function_aliases=True
            ),
        )
        self.assertAllClose(
            ref_output, revived_model.function_aliases["call_alias"](ref_input)
        )
        # Test with a different batch size
        revived_model.function_aliases["call_alias"](tf.random.normal((6, 10)))
//  [python]
def test_low_level_model_export_with_alias(self):
        temp_filepath = os.path.join(self.get_temp_dir(), "exported_model")

        model = get_model()
        ref_input = tf.random.normal((3, 10))
        ref_output = model(ref_input)

        export_archive = export_lib.ExportArchive()
        export_archive.track(model)
        fn = export_archive.add_endpoint(
            "call",
            model.__call__,
            input_signature=[tf.TensorSpec(shape=(None, 10), dtype=tf.float32)],
        )
        export_archive.write_out(
            temp_filepath,
            tf.saved_model.SaveOptions(function_aliases={"call_alias": fn}),
        )
        revived_model = tf.saved_model.load(
            temp_filepath,
            options=tf.saved_model.LoadOptions(
                experimental_load_function_aliases=True
            ),
        )
        self.assertAllClose(
            ref_output, revived_model.function_aliases["call_alias"](ref_input)
        )
        # Test with a different batch size
        revived_model.function_aliases["call_alias"](tf.random.normal((6, 10)))
//  [python]
def test_low_level_model_export_with_dynamic_dims(self, model_type):

        class ReductionLayer(layers.Layer):
            def call(self, inputs):
                return ops.max(inputs, axis=1)

        temp_filepath = os.path.join(self.get_temp_dir(), "exported_model")

        model = get_model(
            model_type,
            input_shape=[(None,), (None,)],
            layer_list=[layers.Concatenate(), ReductionLayer()],
        )
        ref_input = [tf.random.normal((3, 8)), tf.random.normal((3, 6))]
        ref_output = model(ref_input)

        export_archive = export_lib.ExportArchive()
        export_archive.track(model)
        export_archive.add_endpoint(
            "call",
            model.__call__,
            input_signature=[
                [
                    tf.TensorSpec(shape=(None, None), dtype=tf.float32),
                    tf.TensorSpec(shape=(None, None), dtype=tf.float32),
                ]
            ],
        )
        export_archive.write_out(temp_filepath)
        revived_model = tf.compat.v1.saved_model.load(temp_filepath)
        self.assertAllClose(ref_output, revived_model.call(ref_input))
        # Test with a different batch size
        revived_model.call([tf.random.normal((6, 8)), tf.random.normal((6, 6))])
        # Test with a different batch size and different dynamic sizes
        revived_model.call([tf.random.normal((6, 3)), tf.random.normal((6, 5))])
//  [python]
def test_low_level_model_export_with_dynamic_dims(self, model_type):

        class ReductionLayer(layers.Layer):
            def call(self, inputs):
                return ops.max(inputs, axis=1)

        temp_filepath = os.path.join(self.get_temp_dir(), "exported_model")

        model = get_model(
            model_type,
            input_shape=[(None,), (None,)],
            layer_list=[layers.Concatenate(), ReductionLayer()],
        )
        ref_input = [tf.random.normal((3, 8)), tf.random.normal((3, 6))]
        ref_output = model(ref_input)

        export_archive = export_lib.ExportArchive()
        export_archive.track(model)
        export_archive.add_endpoint(
            "call",
            model.__call__,
            input_signature=[
                [
                    tf.TensorSpec(shape=(None, None), dtype=tf.float32),
                    tf.TensorSpec(shape=(None, None), dtype=tf.float32),
                ]
            ],
        )
        export_archive.write_out(temp_filepath)
        revived_model = tf.saved_model.load(temp_filepath)
        self.assertAllClose(ref_output, revived_model.call(ref_input))
        # Test with a different batch size
        revived_model.call([tf.random.normal((6, 8)), tf.random.normal((6, 6))])
        # Test with a different batch size and different dynamic sizes
        revived_model.call([tf.random.normal((6, 3)), tf.random.normal((6, 5))])
//  [python]
def test_low_level_model_export_with_jax2tf_kwargs(self):
        temp_filepath = os.path.join(self.get_temp_dir(), "exported_model")

        model = get_model()
        ref_input = tf.random.normal((3, 10))
        ref_output = model(ref_input)

        export_archive = export_lib.ExportArchive()
        export_archive.track(model)
        export_archive.add_endpoint(
            "call",
            model.__call__,
            input_signature=[tf.TensorSpec(shape=(None, 10), dtype=tf.float32)],
            jax2tf_kwargs={
                "native_serialization": True,
                "native_serialization_platforms": ("cpu", "tpu"),
            },
        )
        with self.assertRaisesRegex(
            ValueError, "native_serialization_platforms.*bogus"
        ):
            export_archive.add_endpoint(
                "call2",
                model.__call__,
                input_signature=[
                    tf.TensorSpec(shape=(None, 10), dtype=tf.float32)
                ],
                jax2tf_kwargs={
                    "native_serialization": True,
                    "native_serialization_platforms": ("cpu", "bogus"),
                },
            )
        export_archive.write_out(temp_filepath)
        revived_model = tf.compat.v1.saved_model.load(temp_filepath)
        self.assertAllClose(ref_output, revived_model.call(ref_input))
//  [python]
def test_low_level_model_export_with_jax2tf_kwargs(self):
        temp_filepath = os.path.join(self.get_temp_dir(), "exported_model")

        model = get_model()
        ref_input = tf.random.normal((3, 10))
        ref_output = model(ref_input)

        export_archive = export_lib.ExportArchive()
        export_archive.track(model)
        export_archive.add_endpoint(
            "call",
            model.__call__,
            input_signature=[tf.TensorSpec(shape=(None, 10), dtype=tf.float32)],
            jax2tf_kwargs={
                "native_serialization": True,
                "native_serialization_platforms": ("cpu", "tpu"),
            },
        )
        with self.assertRaisesRegex(
            ValueError, "native_serialization_platforms.*bogus"
        ):
            export_archive.add_endpoint(
                "call2",
                model.__call__,
                input_signature=[
                    tf.TensorSpec(shape=(None, 10), dtype=tf.float32)
                ],
                jax2tf_kwargs={
                    "native_serialization": True,
                    "native_serialization_platforms": ("cpu", "bogus"),
                },
            )
        export_archive.write_out(temp_filepath)
        revived_model = tf.saved_model.load(temp_filepath)
        self.assertAllClose(ref_output, revived_model.call(ref_input))
//  [python]
def test_layer_export(self):
        temp_filepath = os.path.join(self.get_temp_dir(), "exported_layer")

        layer = layers.BatchNormalization()
        ref_input = tf.random.normal((3, 10))
        ref_output = layer(ref_input)  # Build layer (important)

        export_archive = export_lib.ExportArchive()
        export_archive.track(layer)
        export_archive.add_endpoint(
            "call",
            layer.call,
            input_signature=[tf.TensorSpec(shape=(None, 10), dtype=tf.float32)],
        )
        export_archive.write_out(temp_filepath)
        revived_layer = tf.saved_model.load(temp_filepath)
        self.assertAllClose(ref_output, revived_layer.call(ref_input))
//  [python]
def test_low_level_model_export_with_jax2tf_polymorphic_shapes(self):

        class SquareLayer(layers.Layer):
            def call(self, inputs):
                return ops.matmul(inputs, inputs)

        temp_filepath = os.path.join(self.get_temp_dir(), "exported_model")

        model = CustomModel([SquareLayer()])
        ref_input = tf.random.normal((3, 10, 10))
        ref_output = model(ref_input)
        signature = [tf.TensorSpec(shape=(None, None, None), dtype=tf.float32)]

        with self.assertRaises(TypeError):
            # This will fail because the polymorphic_shapes that is
            # automatically generated will not account for the fact that
            # dynamic dimensions 1 and 2 must have the same value.
            export_archive = export_lib.ExportArchive()
            export_archive.track(model)
            export_archive.add_endpoint(
                "call",
                model.__call__,
                input_signature=signature,
                jax2tf_kwargs={},
            )
            export_archive.write_out(temp_filepath)

        export_archive = export_lib.ExportArchive()
        export_archive.track(model)
        export_archive.add_endpoint(
            "call",
            model.__call__,
            input_signature=signature,
            jax2tf_kwargs={"polymorphic_shapes": ["(batch, a, a)"]},
        )
        export_archive.write_out(temp_filepath)
        revived_model = tf.compat.v1.saved_model.load(temp_filepath)
        self.assertAllClose(ref_output, revived_model.call(ref_input))
//  [python]
def test_low_level_model_export_with_jax2tf_polymorphic_shapes(self):

        class SquareLayer(layers.Layer):
            def call(self, inputs):
                return ops.matmul(inputs, inputs)

        temp_filepath = os.path.join(self.get_temp_dir(), "exported_model")

        model = CustomModel([SquareLayer()])
        ref_input = tf.random.normal((3, 10, 10))
        ref_output = model(ref_input)
        signature = [tf.TensorSpec(shape=(None, None, None), dtype=tf.float32)]

        with self.assertRaises(TypeError):
            # This will fail because the polymorphic_shapes that is
            # automatically generated will not account for the fact that
            # dynamic dimensions 1 and 2 must have the same value.
            export_archive = export_lib.ExportArchive()
            export_archive.track(model)
            export_archive.add_endpoint(
                "call",
                model.__call__,
                input_signature=signature,
                jax2tf_kwargs={},
            )
            export_archive.write_out(temp_filepath)

        export_archive = export_lib.ExportArchive()
        export_archive.track(model)
        export_archive.add_endpoint(
            "call",
            model.__call__,
            input_signature=signature,
            jax2tf_kwargs={"polymorphic_shapes": ["(batch, a, a)"]},
        )
        export_archive.write_out(temp_filepath)
        revived_model = tf.saved_model.load(temp_filepath)
        self.assertAllClose(ref_output, revived_model.call(ref_input))
//  [python]
def test_endpoint_registration_tf_function(self):
        temp_filepath = os.path.join(self.get_temp_dir(), "exported_model")

        model = get_model()
        ref_input = tf.random.normal((3, 10))
        ref_output = model(ref_input)

        # Test variable tracking
        export_archive = export_lib.ExportArchive()
        export_archive.track(model)
        self.assertLen(export_archive.variables, 8)
        self.assertLen(export_archive.trainable_variables, 6)
        self.assertLen(export_archive.non_trainable_variables, 2)

        @tf.function()
        def my_endpoint(x):
            return model(x)

        # Test registering an endpoint that is a tf.function (called)
        my_endpoint(ref_input)  # Trace fn

        export_archive.add_endpoint(
            "call",
            my_endpoint,
        )
        export_archive.write_out(temp_filepath)

        revived_model = tf.compat.v1.saved_model.load(temp_filepath)
        self.assertFalse(hasattr(revived_model, "_tracked"))
        self.assertAllClose(ref_output, revived_model.call(ref_input))
        self.assertLen(revived_model.variables, 8)
        self.assertLen(revived_model.trainable_variables, 6)
        self.assertLen(revived_model.non_trainable_variables, 2)
//  [python]
def test_endpoint_registration_tf_function(self):
        temp_filepath = os.path.join(self.get_temp_dir(), "exported_model")

        model = get_model()
        ref_input = tf.random.normal((3, 10))
        ref_output = model(ref_input)

        # Test variable tracking
        export_archive = export_lib.ExportArchive()
        export_archive.track(model)
        self.assertLen(export_archive.variables, 8)
        self.assertLen(export_archive.trainable_variables, 6)
        self.assertLen(export_archive.non_trainable_variables, 2)

        @tf.function()
        def my_endpoint(x):
            return model(x)

        # Test registering an endpoint that is a tf.function (called)
        my_endpoint(ref_input)  # Trace fn

        export_archive.add_endpoint(
            "call",
            my_endpoint,
        )
        export_archive.write_out(temp_filepath)

        revived_model = tf.saved_model.load(temp_filepath)
        self.assertFalse(hasattr(revived_model, "_tracked"))
        self.assertAllClose(ref_output, revived_model.call(ref_input))
        self.assertLen(revived_model.variables, 8)
        self.assertLen(revived_model.trainable_variables, 6)
        self.assertLen(revived_model.non_trainable_variables, 2)
//  [python]
def test_track_multiple_layers(self):
        temp_filepath = os.path.join(self.get_temp_dir(), "exported_model")
        layer_1 = layers.Dense(2)
        ref_input_1 = tf.random.normal((3, 4))
        ref_output_1 = layer_1(ref_input_1)
        layer_2 = layers.Dense(3)
        ref_input_2 = tf.random.normal((3, 5))
        ref_output_2 = layer_2(ref_input_2)

        export_archive = export_lib.ExportArchive()
        export_archive.add_endpoint(
            "call_1",
            layer_1.call,
            input_signature=[tf.TensorSpec(shape=(None, 4), dtype=tf.float32)],
        )
        export_archive.add_endpoint(
            "call_2",
            layer_2.call,
            input_signature=[tf.TensorSpec(shape=(None, 5), dtype=tf.float32)],
        )
        export_archive.write_out(temp_filepath)
        revived_layer = tf.saved_model.load(temp_filepath)
        self.assertAllClose(ref_output_1, revived_layer.call_1(ref_input_1))
        self.assertAllClose(ref_output_2, revived_layer.call_2(ref_input_2))
//  [python]
def test_non_standard_layer_signature(self):
        temp_filepath = os.path.join(self.get_temp_dir(), "exported_layer")

        layer = layers.MultiHeadAttention(2, 2)
        x1 = tf.random.normal((3, 2, 2))
        x2 = tf.random.normal((3, 2, 2))
        ref_output = layer(x1, x2)  # Build layer (important)
        export_archive = export_lib.ExportArchive()
        export_archive.track(layer)
        export_archive.add_endpoint(
            "call",
            layer.call,
            input_signature=[
                tf.TensorSpec(shape=(None, 2, 2), dtype=tf.float32),
                tf.TensorSpec(shape=(None, 2, 2), dtype=tf.float32),
            ],
        )
        export_archive.write_out(temp_filepath)
        revived_layer = tf.compat.v1.saved_model.load(temp_filepath)
        self.assertAllClose(ref_output, revived_layer.call(x1, x2))
//  [python]
def test_non_standard_layer_signature(self):
        temp_filepath = os.path.join(self.get_temp_dir(), "exported_layer")

        layer = layers.MultiHeadAttention(2, 2)
        x1 = tf.random.normal((3, 2, 2))
        x2 = tf.random.normal((3, 2, 2))
        ref_output = layer(x1, x2)  # Build layer (important)
        export_archive = export_lib.ExportArchive()
        export_archive.track(layer)
        export_archive.add_endpoint(
            "call",
            layer.call,
            input_signature=[
                tf.TensorSpec(shape=(None, 2, 2), dtype=tf.float32),
                tf.TensorSpec(shape=(None, 2, 2), dtype=tf.float32),
            ],
        )
        export_archive.write_out(temp_filepath)
        revived_layer = tf.saved_model.load(temp_filepath)
        self.assertAllClose(ref_output, revived_layer.call(x1, x2))
//  [python]
def test_jax_endpoint_registration_tf_function(self):
        model = get_model()
        ref_input = np.random.normal(size=(3, 10))
        model(ref_input)

        # build a JAX function
        def model_call(x):
            return model(x)

        from jax import default_backend as jax_device
        from jax.experimental import jax2tf

        native_jax_compatible = not (
            jax_device() == "gpu"
            and len(tf.config.list_physical_devices("GPU")) == 0
        )
        # now, convert JAX function
        converted_model_call = jax2tf.convert(
            model_call,
            native_serialization=native_jax_compatible,
            polymorphic_shapes=["(b, 10)"],
        )

        # you can now build a TF inference function
        @tf.function(
            input_signature=[tf.TensorSpec(shape=(None, 10), dtype=tf.float32)],
            autograph=False,
        )
        def infer_fn(x):
            return converted_model_call(x)

        ref_output = infer_fn(ref_input)

        # Export with TF inference function as endpoint
        temp_filepath = os.path.join(self.get_temp_dir(), "my_model")
        export_archive = export_lib.ExportArchive()
        export_archive.track(model)
        export_archive.add_endpoint("serve", infer_fn)
        export_archive.write_out(temp_filepath)

        # Reload and verify outputs
        revived_model = tf.compat.v1.saved_model.load(temp_filepath)
        self.assertFalse(hasattr(revived_model, "_tracked"))
        self.assertAllClose(ref_output, revived_model.serve(ref_input))
        self.assertLen(revived_model.variables, 8)
        self.assertLen(revived_model.trainable_variables, 6)
        self.assertLen(revived_model.non_trainable_variables, 2)

        # Assert all variables wrapped as `tf.Variable`
        assert isinstance(export_archive.variables[0], tf.Variable)
        assert isinstance(export_archive.trainable_variables[0], tf.Variable)
        assert isinstance(
            export_archive.non_trainable_variables[0], tf.Variable
        )
//  [python]
def test_jax_endpoint_registration_tf_function(self):
        model = get_model()
        ref_input = np.random.normal(size=(3, 10))
        model(ref_input)

        # build a JAX function
        def model_call(x):
            return model(x)

        from jax import default_backend as jax_device
        from jax.experimental import jax2tf

        native_jax_compatible = not (
            jax_device() == "gpu"
            and len(tf.config.list_physical_devices("GPU")) == 0
        )
        # now, convert JAX function
        converted_model_call = jax2tf.convert(
            model_call,
            native_serialization=native_jax_compatible,
            polymorphic_shapes=["(b, 10)"],
        )

        # you can now build a TF inference function
        @tf.function(
            input_signature=[tf.TensorSpec(shape=(None, 10), dtype=tf.float32)],
            autograph=False,
        )
        def infer_fn(x):
            return converted_model_call(x)

        ref_output = infer_fn(ref_input)

        # Export with TF inference function as endpoint
        temp_filepath = os.path.join(self.get_temp_dir(), "my_model")
        export_archive = export_lib.ExportArchive()
        export_archive.track(model)
        export_archive.add_endpoint("serve", infer_fn)
        export_archive.write_out(temp_filepath)

        # Reload and verify outputs
        revived_model = tf.saved_model.load(temp_filepath)
        self.assertFalse(hasattr(revived_model, "_tracked"))
        self.assertAllClose(ref_output, revived_model.serve(ref_input))
        self.assertLen(revived_model.variables, 8)
        self.assertLen(revived_model.trainable_variables, 6)
        self.assertLen(revived_model.non_trainable_variables, 2)

        # Assert all variables wrapped as `tf.Variable`
        assert isinstance(export_archive.variables[0], tf.Variable)
        assert isinstance(export_archive.trainable_variables[0], tf.Variable)
        assert isinstance(
            export_archive.non_trainable_variables[0], tf.Variable
        )
//  [python]
def test_model_export_method(self, model_type):
        temp_filepath = os.path.join(self.get_temp_dir(), "exported_model")
        model = get_model(model_type)
        ref_input = tf.random.normal((3, 10))
        ref_output = model(ref_input)

        model.export(temp_filepath)
        revived_model = tf.compat.v1.saved_model.load(temp_filepath)
        self.assertAllClose(ref_output, revived_model.serve(ref_input))
        # Test with a different batch size
        revived_model.serve(tf.random.normal((6, 10)))
//  [python]
def test_model_export_method(self, model_type):
        temp_filepath = os.path.join(self.get_temp_dir(), "exported_model")
        model = get_model(model_type)
        ref_input = tf.random.normal((3, 10))
        ref_output = model(ref_input)

        model.export(temp_filepath)
        revived_model = tf.saved_model.load(temp_filepath)
        self.assertAllClose(ref_output, revived_model.serve(ref_input))
        # Test with a different batch size
        revived_model.serve(tf.random.normal((6, 10)))
//  [python]
def test_variable_collection(self):
        temp_filepath = os.path.join(self.get_temp_dir(), "exported_model")

        model = models.Sequential(
            [
                layers.Input((10,)),
                layers.Dense(2),
                layers.Dense(2),
            ]
        )

        # Test variable tracking
        export_archive = export_lib.ExportArchive()
        export_archive.track(model)
        export_archive.add_endpoint(
            "call",
            model.__call__,
            input_signature=[tf.TensorSpec(shape=(None, 10), dtype=tf.float32)],
        )
        export_archive.add_variable_collection(
            "my_vars", model.layers[1].weights
        )

        self.assertLen(export_archive._tf_trackable.my_vars, 2)
        export_archive.write_out(temp_filepath)
        revived_model = tf.compat.v1.saved_model.load(temp_filepath)
        self.assertLen(revived_model.my_vars, 2)
//  [python]
def test_variable_collection(self):
        temp_filepath = os.path.join(self.get_temp_dir(), "exported_model")

        model = models.Sequential(
            [
                layers.Input((10,)),
                layers.Dense(2),
                layers.Dense(2),
            ]
        )

        # Test variable tracking
        export_archive = export_lib.ExportArchive()
        export_archive.track(model)
        export_archive.add_endpoint(
            "call",
            model.__call__,
            input_signature=[tf.TensorSpec(shape=(None, 10), dtype=tf.float32)],
        )
        export_archive.add_variable_collection(
            "my_vars", model.layers[1].weights
        )

        self.assertLen(export_archive._tf_trackable.my_vars, 2)
        export_archive.write_out(temp_filepath)
        revived_model = tf.saved_model.load(temp_filepath)
        self.assertLen(revived_model.my_vars, 2)
//  [python]
def test_non_standard_layer_signature_with_kwargs(self):
        temp_filepath = os.path.join(self.get_temp_dir(), "exported_layer")

        layer = layers.MultiHeadAttention(2, 2)
        x1 = tf.random.normal((3, 2, 2))
        x2 = tf.random.normal((3, 2, 2))
        ref_output = layer(x1, x2)  # Build layer (important)
        export_archive = export_lib.ExportArchive()
        export_archive.track(layer)
        export_archive.add_endpoint(
            "call",
            layer.call,
            input_signature=[
                tf.TensorSpec(shape=(None, 2, 2), dtype=tf.float32),
                tf.TensorSpec(shape=(None, 2, 2), dtype=tf.float32),
            ],
        )
        export_archive.write_out(temp_filepath)
        revived_layer = tf.compat.v1.saved_model.load(temp_filepath)
        self.assertAllClose(ref_output, revived_layer.call(query=x1, value=x2))
        # Test with a different batch size
        revived_layer.call(
            query=tf.random.normal((6, 2, 2)), value=tf.random.normal((6, 2, 2))
        )
//  [python]
def test_non_standard_layer_signature_with_kwargs(self):
        temp_filepath = os.path.join(self.get_temp_dir(), "exported_layer")

        layer = layers.MultiHeadAttention(2, 2)
        x1 = tf.random.normal((3, 2, 2))
        x2 = tf.random.normal((3, 2, 2))
        ref_output = layer(x1, x2)  # Build layer (important)
        export_archive = export_lib.ExportArchive()
        export_archive.track(layer)
        export_archive.add_endpoint(
            "call",
            layer.call,
            input_signature=[
                tf.TensorSpec(shape=(None, 2, 2), dtype=tf.float32),
                tf.TensorSpec(shape=(None, 2, 2), dtype=tf.float32),
            ],
        )
        export_archive.write_out(temp_filepath)
        revived_layer = tf.saved_model.load(temp_filepath)
        self.assertAllClose(ref_output, revived_layer.call(query=x1, value=x2))
        # Test with a different batch size
        revived_layer.call(
            query=tf.random.normal((6, 2, 2)), value=tf.random.normal((6, 2, 2))
        )
//  [python]
def _convert_labeled_classes_to_k_hot(groundtruth_labeled_classes, num_classes,
                                      map_empty_to_ones=False):
  """Returns k-hot encoding of the labeled classes.

  If map_empty_to_ones is enabled and the input labeled_classes is empty,
  this function assumes all classes are exhaustively labeled, thus returning
  an all-one encoding.

  Args:
    groundtruth_labeled_classes: a Tensor holding a sparse representation of
      labeled classes.
    num_classes: an integer representing the number of classes
    map_empty_to_ones: boolean (default: False).  Set this to be True to default
    to an all-ones result if given an empty `groundtruth_labeled_classes`.
  Returns:
    A k-hot (and 0-indexed) tensor representation of
    `groundtruth_labeled_classes`.
  """

  # If the input labeled_classes is empty, it assumes all classes are
  # exhaustively labeled, thus returning an all-one encoding.
  def true_fn():
    return tf.sparse_to_dense(
        groundtruth_labeled_classes - _LABEL_OFFSET, [num_classes],
        tf.constant(1, dtype=tf.float32),
        validate_indices=False)

  def false_fn():
    return tf.ones(num_classes, dtype=tf.float32)

  if map_empty_to_ones:
    return tf.cond(tf.size(groundtruth_labeled_classes) > 0, true_fn, false_fn)
  return true_fn()
//  [python]
def _convert_labeled_classes_to_k_hot(groundtruth_labeled_classes, num_classes,
                                      map_empty_to_ones=False):
  """Returns k-hot encoding of the labeled classes.

  If map_empty_to_ones is enabled and the input labeled_classes is empty,
  this function assumes all classes are exhaustively labeled, thus returning
  an all-one encoding.

  Args:
    groundtruth_labeled_classes: a Tensor holding a sparse representation of
      labeled classes.
    num_classes: an integer representing the number of classes
    map_empty_to_ones: boolean (default: False).  Set this to be True to default
    to an all-ones result if given an empty `groundtruth_labeled_classes`.
  Returns:
    A k-hot (and 0-indexed) tensor representation of
    `groundtruth_labeled_classes`.
  """

  # If the input labeled_classes is empty, it assumes all classes are
  # exhaustively labeled, thus returning an all-one encoding.
  def true_fn():
    return tf.sparse.to_dense(
        groundtruth_labeled_classes - _LABEL_OFFSET, [num_classes],
        tf.constant(1, dtype=tf.float32),
        validate_indices=False)

  def false_fn():
    return tf.ones(num_classes, dtype=tf.float32)

  if map_empty_to_ones:
    return tf.cond(tf.size(groundtruth_labeled_classes) > 0, true_fn, false_fn)
  return true_fn()
//  [python]
def test_multi_input_output_functional_model(self):
        temp_filepath = os.path.join(self.get_temp_dir(), "exported_model")
        x1 = layers.Input((2,))
        x2 = layers.Input((2,))
        y1 = layers.Dense(3)(x1)
        y2 = layers.Dense(3)(x2)
        model = models.Model([x1, x2], [y1, y2])

        ref_inputs = [tf.random.normal((3, 2)), tf.random.normal((3, 2))]
        ref_outputs = model(ref_inputs)

        export_archive = export_lib.ExportArchive()
        export_archive.track(model)
        export_archive.add_endpoint(
            "serve",
            model.__call__,
            input_signature=[
                [
                    tf.TensorSpec(shape=(None, 2), dtype=tf.float32),
                    tf.TensorSpec(shape=(None, 2), dtype=tf.float32),
                ]
            ],
        )
        export_archive.write_out(temp_filepath)
        revived_model = tf.saved_model.load(temp_filepath)
        self.assertAllClose(ref_outputs[0], revived_model.serve(ref_inputs)[0])
        self.assertAllClose(ref_outputs[1], revived_model.serve(ref_inputs)[1])
        # Test with a different batch size
        revived_model.serve(
            [tf.random.normal((6, 2)), tf.random.normal((6, 2))]
        )

        # Now test dict inputs
        model = models.Model({"x1": x1, "x2": x2}, [y1, y2])

        ref_inputs = {
            "x1": tf.random.normal((3, 2)),
            "x2": tf.random.normal((3, 2)),
        }
        ref_outputs = model(ref_inputs)

        export_archive = export_lib.ExportArchive()
        export_archive.track(model)
        export_archive.add_endpoint(
            "serve",
            model.__call__,
            input_signature=[
                {
                    "x1": tf.TensorSpec(shape=(None, 2), dtype=tf.float32),
                    "x2": tf.TensorSpec(shape=(None, 2), dtype=tf.float32),
                }
            ],
        )
        export_archive.write_out(temp_filepath)
        revived_model = tf.saved_model.load(temp_filepath)
        self.assertAllClose(ref_outputs[0], revived_model.serve(ref_inputs)[0])
        self.assertAllClose(ref_outputs[1], revived_model.serve(ref_inputs)[1])
        # Test with a different batch size
        revived_model.serve(
            {
                "x1": tf.random.normal((6, 2)),
                "x2": tf.random.normal((6, 2)),
            }
        )
//  [python]
def parse_sequence(self, sequence, max_length, delimiter="", default_value="0"):
        """
        split the sequence and convert to dense tensor
        """
        split_sequence = tf.string_split(sequence, delimiter=delimiter)
        split_sequence = tf.sparse_to_dense(sparse_indices=split_sequence.indices,
                                            output_shape=[self.feat_config["batch_size"],
                                                          max_length],
                                            sparse_values=split_sequence.values, default_value=default_value)

        return split_sequence
//  [python]
def parse_sequence(self, sequence, max_length, delimiter="", default_value="0"):
        """
        split the sequence and convert to dense tensor
        """
        split_sequence = tf.string_split(sequence, delimiter=delimiter)
        split_sequence = tf.sparse.to_dense(sparse_indices=split_sequence.indices,
                                            output_shape=[self.feat_config["batch_size"],
                                                          max_length],
                                            sparse_values=split_sequence.values, default_value=default_value)

        return split_sequence
//  [python]
def _test_sparse_to_dense(sparse_indices, sparse_values, default_value, output_shape):
    with tf.Graph().as_default():
        indices = tf.placeholder(
            shape=sparse_indices.shape, dtype=str(sparse_indices.dtype), name="indices"
        )
        values = tf.placeholder(
            shape=sparse_values.shape, dtype=str(sparse_values.dtype), name="values"
        )
        oshape = tf.constant(output_shape, shape=output_shape.shape, dtype=str(output_shape.dtype))

        # Output shape depends on a dynamic input, use VM.
        if default_value is None:
            output = tf.sparse_to_dense(indices, oshape, values)
            compare_tf_with_tvm(
                [sparse_indices, sparse_values], ["indices:0", "values:0"], output.name, mode="vm"
            )
        else:
            dv = tf.placeholder(shape=(), dtype=str(default_value.dtype), name="default_value")
            output = tf.sparse_to_dense(indices, oshape, values, dv)
            compare_tf_with_tvm(
                [sparse_indices, sparse_values, default_value],
                ["indices:0", "values:0", "default_value:0"],
                output.name,
                mode="vm",
            )
//  [python]
def _test_sparse_to_dense(sparse_indices, sparse_values, default_value, output_shape):
    with tf.Graph().as_default():
        indices = tf.placeholder(
            shape=sparse_indices.shape, dtype=str(sparse_indices.dtype), name="indices"
        )
        values = tf.placeholder(
            shape=sparse_values.shape, dtype=str(sparse_values.dtype), name="values"
        )
        oshape = tf.constant(output_shape, shape=output_shape.shape, dtype=str(output_shape.dtype))

        # Output shape depends on a dynamic input, use VM.
        if default_value is None:
            output = tf.sparse.to_dense(tf.sparse.SparseTensor(indices, values, oshape))
            compare_tf_with_tvm(
                [sparse_indices, sparse_values], ["indices:0", "values:0"], output.name, mode="vm"
            )
        else:
            dv = tf.placeholder(shape=(), dtype=str(default_value.dtype), name="default_value")
            output = tf.sparse.to_dense(tf.sparse.SparseTensor(indices, values, oshape), default_value=dv)
            compare_tf_with_tvm(
                [sparse_indices, sparse_values, default_value],
                ["indices:0", "values:0", "default_value:0"],
                output.name,
                mode="vm",
            )
//  [python]
def _single_token_mask(inputs, tgt_len, num_predict, exclude_mask=None):
  """Sample individual tokens as prediction targets."""
  func_mask = tf.equal(inputs, FLAGS.cls_id)
  func_mask = tf.logical_or(func_mask, tf.equal(inputs, FLAGS.sep_id))
  func_mask = tf.logical_or(func_mask, tf.equal(inputs, FLAGS.pad_id))
  if exclude_mask is None:
    exclude_mask = func_mask
  else:
    exclude_mask = tf.logical_or(func_mask, exclude_mask)
  candidate_mask = tf.logical_not(exclude_mask)

  all_indices = tf.range(tgt_len, dtype=tf.int64)
  candidate_indices = tf.boolean_mask(all_indices, candidate_mask)
  masked_pos = tf.random.shuffle(candidate_indices)
  masked_pos = tf.sort(masked_pos[:num_predict])
  target_mask = tf.sparse_to_dense(
      sparse_indices=masked_pos,
      output_shape=[tgt_len],
      sparse_values=1.0,
      default_value=0.0)
  is_target = tf.cast(target_mask, tf.bool)

  return is_target, target_mask
//  [python]
def _single_token_mask(inputs, tgt_len, num_predict, exclude_mask=None):
  """Sample individual tokens as prediction targets."""
  func_mask = tf.equal(inputs, FLAGS.cls_id)
  func_mask = tf.logical_or(func_mask, tf.equal(inputs, FLAGS.sep_id))
  func_mask = tf.logical_or(func_mask, tf.equal(inputs, FLAGS.pad_id))
  if exclude_mask is None:
    exclude_mask = func_mask
  else:
    exclude_mask = tf.logical_or(func_mask, exclude_mask)
  candidate_mask = tf.logical_not(exclude_mask)

  all_indices = tf.range(tgt_len, dtype=tf.int64)
  candidate_indices = tf.boolean_mask(all_indices, candidate_mask)
  masked_pos = tf.random.shuffle(candidate_indices)
  masked_pos = tf.sort(masked_pos[:num_predict])
  target_mask = tf.sparse.to_dense(
      sparse_indices=masked_pos,
      output_shape=[tgt_len],
      sparse_values=1.0,
      default_value=0.0)
  is_target = tf.cast(target_mask, tf.bool)

  return is_target, target_mask
//  [python]
def _convert_labeled_classes_to_k_hot(groundtruth_labeled_classes, num_classes):
  """Returns k-hot encoding of the labeled classes."""

  # If the input labeled_classes is empty, it assumes all classes are
  # exhaustively labeled, thus returning an all-one encoding.
  def true_fn():
    return tf.sparse_to_dense(
        groundtruth_labeled_classes - _LABEL_OFFSET, [num_classes],
        tf.constant(1, dtype=tf.float32),
        validate_indices=False)

  def false_fn():
    return tf.ones(num_classes, dtype=tf.float32)

  return tf.cond(tf.size(groundtruth_labeled_classes) > 0, true_fn, false_fn)
//  [python]
def _convert_labeled_classes_to_k_hot(groundtruth_labeled_classes, num_classes):
  """Returns k-hot encoding of the labeled classes."""

  # If the input labeled_classes is empty, it assumes all classes are
  # exhaustively labeled, thus returning an all-one encoding.
  def true_fn():
    return tf.sparse.to_dense(
        groundtruth_labeled_classes - _LABEL_OFFSET, [num_classes],
        tf.constant(1, dtype=tf.float32),
        validate_indices=False)

  def false_fn():
    return tf.ones(num_classes, dtype=tf.float32)

  return tf.cond(tf.size(groundtruth_labeled_classes) > 0, true_fn, false_fn)