jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def _load_pyfunc(path):
    """
    Load PyFunc implementation. Called by ``pyfunc.load_pyfunc``. This function loads an MLflow
    model with the TensorFlow flavor into a new TensorFlow graph and exposes it behind the
    ``pyfunc.predict`` interface.

    :param path: Local filesystem path to the MLflow Model with the ``tensorflow`` flavor.
    """
    import tensorflow as t

    (
        tf_saved_model_dir,
        tf_meta_graph_tags,
        tf_signature_def_key,
    ) = _get_and_parse_flavor_configuration(model_path=path)
    if t.__version__ < "2.0.0":
        tf_graph = t.Graph()
        tf_sess = t.Session(graph=tf_graph)
        with tf_graph.as_default():
            signature_def = _load_tensorflow_saved_model(
                tf_saved_model_dir=tf_saved_model_dir,
                tf_sess=tf_sess,
                tf_meta_graph_tags=tf_meta_graph_tags,
                tf_signature_def_key=tf_signature_def_key,
            )

        return _TFWrapper(tf_sess=tf_sess, tf_graph=tf_graph, signature_def=signature_def)
    else:
        tf_graph = t.Graph()
        tf_sess = t.Session(graph=tf_graph)
        with tf_graph.as_default():
            signature_def = _load_tensorflow_saved_model(
                tf_saved_model_dir=tf_saved_model_dir,
                tf_sess=tf_sess,
                tf_meta_graph_tags=tf_meta_graph_tags,
                tf_signature_def_key=tf_signature_def_key,
            )
        return _TF2Wrapper(tf_sess=tf_sess, tf_graph=tf_graph, signature_def=signature_def)
//  [python]
def _load_pyfunc(path):
    """
    Load PyFunc implementation. Called by ``pyfunc.load_pyfunc``. This function loads an MLflow
    model with the TensorFlow flavor into a new TensorFlow graph and exposes it behind the
    ``pyfunc.predict`` interface.

    :param path: Local filesystem path to the MLflow Model with the ``tensorflow`` flavor.
    """
    import tensorflow

    (
        tf_saved_model_dir,
        tf_meta_graph_tags,
        tf_signature_def_key,
    ) = _get_and_parse_flavor_configuration(model_path=path)
    if LooseVersion(tensorflow.__version__) < LooseVersion("2.0.0"):
        tf_graph = tensorflow.Graph()
        tf_sess = tensorflow.Session(graph=tf_graph)
        with tf_graph.as_default():
            signature_def = _load_tensorflow_saved_model(
                tf_saved_model_dir=tf_saved_model_dir,
                tf_sess=tf_sess,
                tf_meta_graph_tags=tf_meta_graph_tags,
                tf_signature_def_key=tf_signature_def_key,
            )

        return _TFWrapper(tf_sess=tf_sess, tf_graph=tf_graph, signature_def=signature_def)
    else:
        loaded_model = tensorflow.saved_model.load(  # pylint: disable=no-value-for-parameter
            export_dir=tf_saved_model_dir, tags=tf_meta_graph_tags
        )
        return _TF2Wrapper(infer=loaded_model.signatures[tf_signature_def_key])
//  [python]
def test_saved_model(self):
        for tf_tokenizer in self.tf_tokenizers:
            model = ModelToSave(tokenizer=tf_tokenizer)
            test_inputs = tf.convert_to_tensor([self.test_sentences[0]])
            out = model.serving(test_inputs)  # Build model with some sample inputs
            with TemporaryDirectory() as tempdir:
                save_path = Path(tempdir) / "saved.model"
                tf.saved_model.save(model, save_path, signatures={"serving_default": model.serving})
                loaded_model = tf.compat.v1.saved_model.load(save_path)
            loaded_output = loaded_model.signatures["serving_default"](test_inputs)["output_0"]
            # We may see small differences because the loaded model is compiled, so we need an epsilon for the test
            self.assertTrue(tf.reduce_all(out == loaded_output))
//  [python]
def test_saved_model(self):
        for tf_tokenizer in self.tf_tokenizers:
            model = ModelToSave(tokenizer=tf_tokenizer)
            test_inputs = tf.convert_to_tensor([self.test_sentences[0]])
            out = model.serving(test_inputs)  # Build model with some sample inputs
            with TemporaryDirectory() as tempdir:
                save_path = Path(tempdir) / "saved.model"
                tf.saved_model.save(model, save_path, signatures={"serving_default": model.serving})
                loaded_model = tf.saved_model.load(save_path)
            loaded_output = loaded_model.signatures["serving_default"](test_inputs)["output_0"]
            # We may see small differences because the loaded model is compiled, so we need an epsilon for the test
            self.assertTrue(tf.reduce_all(out == loaded_output))
//  [python]
def test_call_tf_graph_save_and_load(self):
    def jax_func(x):
      def func_tf(x):
        return tf.math.sin(x)

      return jnp.cos(
          jax2tf.call_tf(func_tf, output_shape_dtype=x, call_tf_graph=True)(x)
      )
    data_inputs = (np.array([0.5, 0.7], dtype=np.float32),)

    def tf_func(the_input):
      res = jax2tf.convert(jax_func, native_serialization=True)(the_input)
      return tf.identity(res, name="the_result")

    jit_tf_func = tf.function(
        tf_func,
        autograph=False,
        jit_compile=True,
    )
    # The next line is necessary to reproduce this issue. It trigger TF
    # ConcreteFunction tracing. Otherwise, you will fail with another error
    # `Found zero restored functions for caller function`.
    _ = jit_tf_func.get_concrete_function(*data_inputs)
    module = tf.Module()
    module.call = jit_tf_func # Switching to concrete_function works.
    root_dir = self.create_tempdir()
    saved_model_dir = os.path.join(root_dir, "saved_model")
    tf.saved_model.save(
        module,
        saved_model_dir,
        options=tf.saved_model.SaveOptions(experimental_custom_gradients=False),
    )
    loaded_model = tf.compat.v1.saved_model.load(saved_model_dir)
    res = loaded_model.call(*data_inputs)
    self.assertAllClose(jax_func(*data_inputs), res)
//  [python]
def test_call_tf_graph_save_and_load(self):
    def jax_func(x):
      def func_tf(x):
        return tf.math.sin(x)

      return jnp.cos(
          jax2tf.call_tf(func_tf, output_shape_dtype=x, call_tf_graph=True)(x)
      )
    data_inputs = (np.array([0.5, 0.7], dtype=np.float32),)

    def tf_func(the_input):
      res = jax2tf.convert(jax_func, native_serialization=True)(the_input)
      return tf.identity(res, name="the_result")

    jit_tf_func = tf.function(
        tf_func,
        autograph=False,
        jit_compile=True,
    )
    # The next line is necessary to reproduce this issue. It trigger TF
    # ConcreteFunction tracing. Otherwise, you will fail with another error
    # `Found zero restored functions for caller function`.
    _ = jit_tf_func.get_concrete_function(*data_inputs)
    module = tf.Module()
    module.call = jit_tf_func # Switching to concrete_function works.
    root_dir = self.create_tempdir()
    saved_model_dir = os.path.join(root_dir, "saved_model")
    tf.saved_model.save(
        module,
        saved_model_dir,
        options=tf.saved_model.SaveOptions(experimental_custom_gradients=False),
    )
    loaded_model = tf.saved_model.load(saved_model_dir)
    res = loaded_model.call(*data_inputs)
    self.assertAllClose(jax_func(*data_inputs), res)
//  [python]
def __init__(self, saved_model=None, serialize_input=True):
    """Initializes BLEURT from a SavedModel."""
    assert saved_model
    self.saved_model_path = saved_model
    self.serialize_input = serialize_input
    logging.info("Reading checkpoint {}.".format(saved_model))
    imported_model = tf.compat.v1.saved_model.load(saved_model)
    self.bleurt_model_ops = imported_model.signatures["serving_default"]
    logging.info("BLEURT initialized.")
//  [python]
def __init__(self, saved_model=None, serialize_input=True):
    """Initializes BLEURT from a SavedModel."""
    assert saved_model
    self.saved_model_path = saved_model
    self.serialize_input = serialize_input
    logging.info("Reading checkpoint {}.".format(saved_model))
    imported_model = tf.saved_model.load(saved_model)
    self.bleurt_model_ops = imported_model.signatures["serving_default"]
    logging.info("BLEURT initialized.")
//  [python]
def initialize(self):
    """Creates the computation graph and the session."""
    logging.info("Loading model.")

    self._bleurt_graph = tf.Graph()
    with self._bleurt_graph.as_default():

      imported = tf.compat.v1.saved_model.load(self.checkpoint)
      bleurt_model_ops = imported.signatures["serving_default"]
      self._bleurt_ops = bleurt_model_ops(
          input_ids=tf.compat.v1.placeholder(tf.int64, name="input_ids"),
          input_mask=tf.compat.v1.placeholder(tf.int64, name="input_mask"),
          segment_ids=tf.compat.v1.placeholder(tf.int64, name="segment_ids"))

      init_op = tf.group(tf.compat.v1.global_variables_initializer(),
                         tf.compat.v1.tables_initializer())

    self.session = tf.compat.v1.Session(graph=self._bleurt_graph)
    self.session.run(init_op)

    logging.info("Done.")
//  [python]
def initialize(self):
    """Creates the computation graph and the session."""
    logging.info("Loading model.")

    self._bleurt_graph = tf.Graph()
    with self._bleurt_graph.as_default():

      imported = tf.saved_model.load(self.checkpoint)
      bleurt_model_ops = imported.signatures["serving_default"]
      self._bleurt_ops = bleurt_model_ops(
          input_ids=tf.compat.v1.placeholder(tf.int64, name="input_ids"),
          input_mask=tf.compat.v1.placeholder(tf.int64, name="input_mask"),
          segment_ids=tf.compat.v1.placeholder(tf.int64, name="segment_ids"))

      init_op = tf.group(tf.compat.v1.global_variables_initializer(),
                         tf.compat.v1.tables_initializer())

    self.session = tf.compat.v1.Session(graph=self._bleurt_graph)
    self.session.run(init_op)

    logging.info("Done.")
//  [python]
def model(self):
        """Return model in AutoTrackable object."""
        if self._auto_trackable:
            return self._auto_trackable

        root = os.path.abspath(os.path.expanduser(DEFAULT_WORKSPACE))
        root += str(time.time())
        if os.path.exists(root):
            shutil.rmtree(root)
        os.makedirs(root, exist_ok=True)
        if not self._sess:
            self._load_sess(self._model, **self.kwargs)
        _, builder = self.build_saved_model(root)
        builder.save()
        model = tf.compat.v1.saved_model.load(root)
        shutil.rmtree(root)
        self._auto_trackable = model
        return model
//  [python]
def model(self):
        """Return model in AutoTrackable object."""
        if self._auto_trackable:
            return self._auto_trackable

        root = os.path.abspath(os.path.expanduser(DEFAULT_WORKSPACE))
        root += str(time.time())
        if os.path.exists(root):
            shutil.rmtree(root)
        os.makedirs(root, exist_ok=True)
        if not self._sess:
            self._load_sess(self._model, **self.kwargs)
        _, builder = self.build_saved_model(root)
        builder.save()
        model = tf.saved_model.load(root)
        shutil.rmtree(root)
        self._auto_trackable = model
        return model
//  [python]
def __init__(self, model_dir, precision):
        self.saved_model_loaded = tf.compat.v1.saved_model.load(model_dir, tags=[tag_constants.SERVING])
        self.graph_func = self.saved_model_loaded.signatures[signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY]
        self.precision = tf.float16 if precision == "amp" else tf.float32
//  [python]
def __init__(self, model_dir, precision):
        self.saved_model_loaded = tf.saved_model.load(model_dir, tags=[tag_constants.SERVING])
        self.graph_func = self.saved_model_loaded.signatures[signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY]
        self.precision = tf.float16 if precision == "amp" else tf.float32
//  [python]
def tf2_benchmark():
  """Run TF2 benchmark and inference."""
  model = build_tf2_model()
  model_cfg = model.cfg
  isize = FLAGS.image_size or model.cfg.eval.isize
  if FLAGS.export_dir:
    tf.saved_model.save(model, FLAGS.export_dir)
    model = tf.compat.v1.saved_model.load(FLAGS.export_dir)

  batch_size = FLAGS.batch_size
  data_dtype = tf.float16 if FLAGS.mixed_precision else tf.float32
  imgs = tf.ones((batch_size, isize, isize, 3), dtype=data_dtype)
  if model_cfg.model.data_format == 'channels_first':
    imgs = tf.transpose(imgs, [0, 3, 1, 2])

  @tf.function
  def f(x):
    return model(x, training=False)

  print('starting warmup.')
  for _ in range(10):  # warmup runs.
    f(imgs)

  print('start benchmark.')
  start = time.perf_counter()
  for _ in range(10):
    f(imgs)
  end = time.perf_counter()
  inference_time = (end - start) / 10

  print('Per batch inference time: ', inference_time)
  print('FPS: ', batch_size / inference_time)
//  [python]
def tf2_benchmark():
  """Run TF2 benchmark and inference."""
  model = build_tf2_model()
  model_cfg = model.cfg
  isize = FLAGS.image_size or model.cfg.eval.isize
  if FLAGS.export_dir:
    tf.saved_model.save(model, FLAGS.export_dir)
    model = tf.saved_model.load(FLAGS.export_dir)

  batch_size = FLAGS.batch_size
  data_dtype = tf.float16 if FLAGS.mixed_precision else tf.float32
  imgs = tf.ones((batch_size, isize, isize, 3), dtype=data_dtype)
  if model_cfg.model.data_format == 'channels_first':
    imgs = tf.transpose(imgs, [0, 3, 1, 2])

  @tf.function
  def f(x):
    return model(x, training=False)

  print('starting warmup.')
  for _ in range(10):  # warmup runs.
    f(imgs)

  print('start benchmark.')
  start = time.perf_counter()
  for _ in range(10):
    f(imgs)
  end = time.perf_counter()
  inference_time = (end - start) / 10

  print('Per batch inference time: ', inference_time)
  print('FPS: ', batch_size / inference_time)
//  [python]
def _load_model(self):
        """Retrieves all the tag-sets available in the SavedModel.
        Args:
          saved_model_dir: Directory containing the SavedModel.
        Returns:
          String representation of all tag-sets in the SavedModel.
        """
        sess = tf.Session()
        MetaGraphDef = tf.saved_model.loader.load(sess, [tf.saved_model.SERVING],
                                        os.path.join(self.dir, 'model'))
        sig = None
        if tf.saved_model.signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY not in MetaGraphDef.signature_def:
            for sig_itr in MetaGraphDef.signature_def:
                if (sig_itr != tf.saved_model.signature_constants.INIT_OP_SIGNATURE_DEF_KEY
                        and sig_itr != tf.saved_model.signature_constants.TRAIN_OP_SIGNATURE_DEF_KEY):
                    sig = sig_itr
                    break
        else:
            sig = tf.saved_model.signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY

        assert sig is not None, "unable to load model , expected " + tf.saved_model.signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY + " signature"

        self.graph = MetaGraphDef.graph_def
        self.signature = MetaGraphDef.signature_def[sig]
//  [python]
def _load_model(self):
        """Retrieves all the tag-sets available in the SavedModel.
        Args:
          saved_model_dir: Directory containing the SavedModel.
        Returns:
          String representation of all tag-sets in the SavedModel.
        """
        sess = tf.Session()
        MetaGraphDef = saved_model.load(sess, [saved_model.SERVING],
                                        os.path.join(self.dir, 'model'))
        sig = None
        if DEFAULT_SERVING_SIGNATURE_DEF_KEY not in MetaGraphDef.signature_def:
            for sig_itr in MetaGraphDef.signature_def:
                if (sig_itr != INIT_OP_SIGNATURE_DEF_KEY
                        and sig_itr != TRAIN_OP_SIGNATURE_DEF_KEY):
                    sig = sig_itr
                    break
        else:
            sig = DEFAULT_SERVING_SIGNATURE_DEF_KEY

        assert sig is not None, "unable to load model , expected " + DEFAULT_SERVING_SIGNATURE_DEF_KEY + " signature"

        self.graph = MetaGraphDef.graph_def
        self.signature = MetaGraphDef.signature_def[sig]
//  [python]
def export_model(model_dir, prec, tf_trt_model_dir=None):
    model = tf.saved_model.load(model_dir)
    input_shape = [1, 384]
    dummy_input = tf.constant(tf.zeros(input_shape, dtype=tf.int32))
    x = [
        tf.constant(dummy_input, name='input_word_ids'),
        tf.constant(dummy_input, name='input_mask'),
        tf.constant(dummy_input, name='input_type_ids'),
    ]
    _ = model(x)

    trt_prec = trt.TrtPrecisionMode.FP32 if prec == "fp32" else trt.TrtPrecisionMode.FP16
    converter = trt.TrtGraphConverterV2(
        input_saved_model_dir=model_dir,
        conversion_params=trt.TrtConversionParams(precision_mode=trt_prec),
    )
    converter.convert()
    tf_trt_model_dir = tf_trt_model_dir or f'/tmp/tf-trt_model_{prec}'
    converter.save(tf_trt_model_dir)
    print(f"TF-TRT model saved at {tf_trt_model_dir}")
//  [python]
def __init__(self, model_dir, precision):
        temp_tftrt_dir = f"/tmp/tf-trt_model_{precision}"
        export_model(model_dir, precision, temp_tftrt_dir)
        saved_model_loaded = tf.compat.v1.saved_model.load(temp_tftrt_dir, tags=[tag_constants.SERVING])
        print(f"TF-TRT model loaded from {temp_tftrt_dir}")
        self.graph_func = saved_model_loaded.signatures[signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY]
        self.precision = tf.float16 if precision == "amp" else tf.float32
//  [python]
def __init__(self, model_dir, precision):
        temp_tftrt_dir = f"/tmp/tf-trt_model_{precision}"
        export_model(model_dir, precision, temp_tftrt_dir)
        saved_model_loaded = tf.saved_model.load(temp_tftrt_dir, tags=[tag_constants.SERVING])
        print(f"TF-TRT model loaded from {temp_tftrt_dir}")
        self.graph_func = saved_model_loaded.signatures[signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY]
        self.precision = tf.float16 if precision == "amp" else tf.float32
//  [python]
def test_sequential(self):
        model = models.Sequential([layers.Dense(1)])
        model.compile(loss="mse", optimizer="adam")
        X_train = np.random.rand(100, 3)
        y_train = np.random.rand(100, 1)
        model.fit(X_train, y_train)
        path = os.path.join(self.get_temp_dir(), "my_keras_model")
        tf.saved_model.save(model, path)
        restored_model = tf.compat.v1.saved_model.load(path)
        self.assertAllClose(
            model(X_train),
            restored_model.signatures["serving_default"](
                tf.convert_to_tensor(X_train, dtype=tf.float32)
            )["output_0"],
            rtol=1e-4,
            atol=1e-4,
        )
//  [python]
def test_sequential(self):
        model = models.Sequential([layers.Dense(1)])
        model.compile(loss="mse", optimizer="adam")
        X_train = np.random.rand(100, 3)
        y_train = np.random.rand(100, 1)
        model.fit(X_train, y_train)
        path = os.path.join(self.get_temp_dir(), "my_keras_model")
        tf.saved_model.save(model, path)
        restored_model = tf.saved_model.load(path)
        self.assertAllClose(
            model(X_train),
            restored_model.signatures["serving_default"](
                tf.convert_to_tensor(X_train, dtype=tf.float32)
            )["output_0"],
            rtol=1e-4,
            atol=1e-4,
        )
//  [python]
def test_functional(self):
        inputs = layers.Input(shape=(3,))
        x = layers.Dense(1, name="first_dense")(inputs)
        outputs = layers.Dense(1, name="second_dense")(x)
        model = models.Model(inputs, outputs)
        model.compile(
            optimizer="adam",
            loss="mse",
        )
        X_train = np.random.rand(100, 3)
        y_train = np.random.rand(100, 1)
        model.fit(X_train, y_train)
        path = os.path.join(self.get_temp_dir(), "my_keras_model")
        tf.saved_model.save(model, path)
        restored_model = tf.compat.v1.saved_model.load(path)
        self.assertAllClose(
            model(X_train),
            restored_model.signatures["serving_default"](
                tf.convert_to_tensor(X_train, dtype=tf.float32)
            )["output_0"],
            rtol=1e-4,
            atol=1e-4,
        )
//  [python]
def test_functional(self):
        inputs = layers.Input(shape=(3,))
        x = layers.Dense(1, name="first_dense")(inputs)
        outputs = layers.Dense(1, name="second_dense")(x)
        model = models.Model(inputs, outputs)
        model.compile(
            optimizer="adam",
            loss="mse",
        )
        X_train = np.random.rand(100, 3)
        y_train = np.random.rand(100, 1)
        model.fit(X_train, y_train)
        path = os.path.join(self.get_temp_dir(), "my_keras_model")
        tf.saved_model.save(model, path)
        restored_model = tf.saved_model.load(path)
        self.assertAllClose(
            model(X_train),
            restored_model.signatures["serving_default"](
                tf.convert_to_tensor(X_train, dtype=tf.float32)
            )["output_0"],
            rtol=1e-4,
            atol=1e-4,
        )
//  [python]
def test_subclassed(self):
        model = CustomModelX()
        model.compile(
            optimizer="adam",
            loss="mse",
            metrics=[metrics.Hinge(), "mse"],
        )
        X_train = np.random.rand(100, 3)
        y_train = np.random.rand(100, 1)
        model.fit(X_train, y_train)
        path = os.path.join(self.get_temp_dir(), "my_keras_model")
        tf.saved_model.save(model, path)
        restored_model = tf.saved_model.load(path)
        self.assertAllClose(
            model(X_train),
            restored_model.signatures["serving_default"](
                tf.convert_to_tensor(X_train, dtype=tf.float32)
            )["output_0"],
            rtol=1e-4,
            atol=1e-4,
        )
//  [python]
def test_custom_model_and_layer(self):
        @object_registration.register_keras_serializable(package="my_package")
        class CustomLayer(layers.Layer):
            def __call__(self, inputs):
                return inputs

        @object_registration.register_keras_serializable(package="my_package")
        class Model(models.Model):
            def __init__(self):
                super().__init__()
                self.layer = CustomLayer()

            @tf.function(input_signature=[tf.TensorSpec([None, 1])])
            def call(self, inputs):
                return self.layer(inputs)

        model = Model()
        inp = np.array([[1.0]])
        result = model(inp)
        path = os.path.join(self.get_temp_dir(), "my_keras_model")
        tf.saved_model.save(model, path)
        restored_model = tf.compat.v1.saved_model.load(path)
        self.assertAllClose(
            result,
            restored_model.call(inp),
            rtol=1e-4,
            atol=1e-4,
        )
//  [python]
def test_custom_model_and_layer(self):
        @object_registration.register_keras_serializable(package="my_package")
        class CustomLayer(layers.Layer):
            def __call__(self, inputs):
                return inputs

        @object_registration.register_keras_serializable(package="my_package")
        class Model(models.Model):
            def __init__(self):
                super().__init__()
                self.layer = CustomLayer()

            @tf.function(input_signature=[tf.TensorSpec([None, 1])])
            def call(self, inputs):
                return self.layer(inputs)

        model = Model()
        inp = np.array([[1.0]])
        result = model(inp)
        path = os.path.join(self.get_temp_dir(), "my_keras_model")
        tf.saved_model.save(model, path)
        restored_model = tf.saved_model.load(path)
        self.assertAllClose(
            result,
            restored_model.call(inp),
            rtol=1e-4,
            atol=1e-4,
        )
//  [python]
def test_multi_input_model(self):
        input_1 = layers.Input(shape=(3,))
        input_2 = layers.Input(shape=(5,))
        model = models.Model([input_1, input_2], [input_1, input_2])
        path = os.path.join(self.get_temp_dir(), "my_keras_model")

        tf.saved_model.save(model, path)
        restored_model = tf.saved_model.load(path)
        input_arr_1 = np.random.random((1, 3)).astype("float32")
        input_arr_2 = np.random.random((1, 5)).astype("float32")

        outputs = restored_model.signatures["serving_default"](
            inputs=tf.convert_to_tensor(input_arr_1, dtype=tf.float32),
            inputs_1=tf.convert_to_tensor(input_arr_2, dtype=tf.float32),
        )

        self.assertAllClose(
            input_arr_1, outputs["output_0"], rtol=1e-4, atol=1e-4
        )
        self.assertAllClose(
            input_arr_2, outputs["output_1"], rtol=1e-4, atol=1e-4
        )
//  [python]
def test_multi_input_custom_model_and_layer(self):
        @object_registration.register_keras_serializable(package="my_package")
        class CustomLayer(layers.Layer):
            def __call__(self, *input_list):
                self.add_loss(input_list[-2] * 2)
                return sum(input_list)

        @object_registration.register_keras_serializable(package="my_package")
        class CustomModel(models.Model):
            def build(self, input_shape):
                super().build(input_shape)
                self.layer = CustomLayer()

            @tf.function
            def call(self, *inputs):
                inputs = list(inputs)
                return self.layer(*inputs)

        model = CustomModel()
        inp = [
            tf.constant(i, shape=[1, 1], dtype=tf.float32) for i in range(1, 4)
        ]
        expected = model(*inp)
        path = os.path.join(self.get_temp_dir(), "my_keras_model")
        tf.saved_model.save(model, path)
        restored_model = tf.saved_model.load(path)
        output = restored_model.call(*inp)
        self.assertAllClose(expected, output, rtol=1e-4, atol=1e-4)
//  [python]
def test_fixed_signature_string_dtype(self):
        @object_registration.register_keras_serializable(package="my_package")
        class Adder(models.Model):
            @tf.function(
                input_signature=[tf.TensorSpec(shape=[], dtype=tf.string)]
            )
            def concat(self, x):
                return x + x

        model = Adder()
        path = os.path.join(self.get_temp_dir(), "my_keras_model")
        tf.saved_model.save(model, path)
        restored_model = tf.compat.v1.saved_model.load(path)
        self.assertEqual(model.concat("hello"), restored_model.concat("hello"))
//  [python]
def test_fixed_signature_string_dtype(self):
        @object_registration.register_keras_serializable(package="my_package")
        class Adder(models.Model):
            @tf.function(
                input_signature=[tf.TensorSpec(shape=[], dtype=tf.string)]
            )
            def concat(self, x):
                return x + x

        model = Adder()
        path = os.path.join(self.get_temp_dir(), "my_keras_model")
        tf.saved_model.save(model, path)
        restored_model = tf.saved_model.load(path)
        self.assertEqual(model.concat("hello"), restored_model.concat("hello"))
//  [python]
def test_list_trackable_children_tracking(self):
        @object_registration.register_keras_serializable(package="my_package")
        class CustomLayerList(layers.Layer):
            def __init__(self):
                super().__init__()
                self.sublayers = [
                    layers.Dense(2),
                    layers.Dense(2),
                ]

            def call(self, inputs):
                x = inputs
                for sublayer in self.sublayers:
                    x = sublayer(x)
                return x

        inputs = layers.Input(shape=(1,))
        outputs = CustomLayerList()(inputs)
        model = models.Model(inputs, outputs)

        inp = np.array([[1.0]])
        expected = model(inp)
        path = os.path.join(self.get_temp_dir(), "my_keras_model")
        tf.saved_model.save(model, path)
        restored_model = tf.compat.v1.saved_model.load(path)
        self.assertAllClose(
            expected,
            restored_model.signatures["serving_default"](
                tf.convert_to_tensor(inp, dtype=tf.float32)
            )["output_0"],
            rtol=1e-4,
            atol=1e-4,
        )
//  [python]
def test_list_trackable_children_tracking(self):
        @object_registration.register_keras_serializable(package="my_package")
        class CustomLayerList(layers.Layer):
            def __init__(self):
                super().__init__()
                self.sublayers = [
                    layers.Dense(2),
                    layers.Dense(2),
                ]

            def call(self, inputs):
                x = inputs
                for sublayer in self.sublayers:
                    x = sublayer(x)
                return x

        inputs = layers.Input(shape=(1,))
        outputs = CustomLayerList()(inputs)
        model = models.Model(inputs, outputs)

        inp = np.array([[1.0]])
        expected = model(inp)
        path = os.path.join(self.get_temp_dir(), "my_keras_model")
        tf.saved_model.save(model, path)
        restored_model = tf.saved_model.load(path)
        self.assertAllClose(
            expected,
            restored_model.signatures["serving_default"](
                tf.convert_to_tensor(inp, dtype=tf.float32)
            )["output_0"],
            rtol=1e-4,
            atol=1e-4,
        )
//  [python]
def test_dict_trackable_children_tracking(self):
        @object_registration.register_keras_serializable(package="my_package")
        class CustomLayerDict(layers.Layer):
            def __init__(self):
                super().__init__()
                self.sublayers = {
                    "first_layer": layers.Dense(2),
                    "second_layer": layers.Dense(2),
                }

            def call(self, inputs):
                x = inputs
                for key, sublayer in self.sublayers.items():
                    x = sublayer(x)
                return x

        inputs = layers.Input(shape=(1,))
        outputs = CustomLayerDict()(inputs)
        model = models.Model(inputs, outputs)

        inp = np.array([[1.0]])
        expected = model(inp)
        path = os.path.join(self.get_temp_dir(), "my_keras_model")
        tf.saved_model.save(model, path)
        restored_model = tf.compat.v1.saved_model.load(path)
        self.assertAllClose(
            expected,
            restored_model.signatures["serving_default"](
                tf.convert_to_tensor(inp, dtype=tf.float32)
            )["output_0"],
            rtol=1e-4,
            atol=1e-4,
        )
//  [python]
def test_dict_trackable_children_tracking(self):
        @object_registration.register_keras_serializable(package="my_package")
        class CustomLayerDict(layers.Layer):
            def __init__(self):
                super().__init__()
                self.sublayers = {
                    "first_layer": layers.Dense(2),
                    "second_layer": layers.Dense(2),
                }

            def call(self, inputs):
                x = inputs
                for key, sublayer in self.sublayers.items():
                    x = sublayer(x)
                return x

        inputs = layers.Input(shape=(1,))
        outputs = CustomLayerDict()(inputs)
        model = models.Model(inputs, outputs)

        inp = np.array([[1.0]])
        expected = model(inp)
        path = os.path.join(self.get_temp_dir(), "my_keras_model")
        tf.saved_model.save(model, path)
        restored_model = tf.saved_model.load(path)
        self.assertAllClose(
            expected,
            restored_model.signatures["serving_default"](
                tf.convert_to_tensor(inp, dtype=tf.float32)
            )["output_0"],
            rtol=1e-4,
            atol=1e-4,
        )
//  [python]
def test_signatures_path(self):
        model = CustomSignatureModel()
        model_with_signature_path = os.path.join(
            self.get_temp_dir(), "model_with_signature"
        )
        call = model.__call__.get_concrete_function(
            tf.TensorSpec(None, tf.float32)
        )

        tf.saved_model.save(model, model_with_signature_path, signatures=call)
        restored_model = tf.compat.v1.saved_model.load(model_with_signature_path)
        self.assertEqual(
            list(restored_model.signatures.keys()), ["serving_default"]
        )
//  [python]
def test_signatures_path(self):
        model = CustomSignatureModel()
        model_with_signature_path = os.path.join(
            self.get_temp_dir(), "model_with_signature"
        )
        call = model.__call__.get_concrete_function(
            tf.TensorSpec(None, tf.float32)
        )

        tf.saved_model.save(model, model_with_signature_path, signatures=call)
        restored_model = tf.saved_model.load(model_with_signature_path)
        self.assertEqual(
            list(restored_model.signatures.keys()), ["serving_default"]
        )
//  [python]
def test_non_fixed_signature_string_dtype(self):
        @object_registration.register_keras_serializable(package="my_package")
        class Adder(models.Model):
            @tf.function
            def concat(self, x):
                return x + x

        model = Adder()

        no_fn_path = os.path.join(self.get_temp_dir(), "my_keras_model_no_fn")
        tf.saved_model.save(model, no_fn_path)
        restored_model = tf.saved_model.load(no_fn_path)
        with self.assertRaisesRegex(ValueError, "zero restored functions"):
            _ = restored_model.concat("hello")

        path = os.path.join(self.get_temp_dir(), "my_keras_model")
        tf.saved_model.save(
            model,
            path,
            signatures=model.concat.get_concrete_function(
                tf.TensorSpec(shape=[], dtype=tf.string, name="string_input")
            ),
        )
        restored_model = tf.saved_model.load(path)
        self.assertEqual(model.concat("hello"), restored_model.concat("hello"))
//  [python]
def test_multiple_signatures_dict_path(self):
        model = CustomSignatureModel()
        model_multiple_signatures_path = os.path.join(
            self.get_temp_dir(), "model_with_multiple_signatures"
        )
        call = model.__call__.get_concrete_function(
            tf.TensorSpec(None, tf.float32)
        )
        signatures = {
            "serving_default": call,
            "array_input": model.__call__.get_concrete_function(
                tf.TensorSpec([None], tf.float32)
            ),
        }

        tf.saved_model.save(
            model, model_multiple_signatures_path, signatures=signatures
        )
        restored_model = tf.saved_model.load(model_multiple_signatures_path)
        self.assertEqual(
            list(restored_model.signatures.keys()),
            ["serving_default", "array_input"],
        )
//  [python]
def test_export_for_inference(self):
        for tf_tokenizer in self.tf_tokenizers:
            model = ModelToSave(tokenizer=tf_tokenizer)
            test_inputs = tf.convert_to_tensor(self.test_sentences)
            out = model(test_inputs)  # Build model with some sample inputs
            with TemporaryDirectory() as tempdir:
                save_path = Path(tempdir) / "saved.model"
                model.export(save_path)
                loaded_model = tf.compat.v1.saved_model.load(save_path)
            loaded_output = loaded_model.serve(test_inputs)
            # We may see small differences because the loaded model is compiled, so we need an epsilon for the test
            self.assertLessEqual(tf.reduce_max(tf.abs(out - loaded_output)), 1e-5)
//  [python]
def test_export_for_inference(self):
        for tf_tokenizer in self.tf_tokenizers:
            model = ModelToSave(tokenizer=tf_tokenizer)
            test_inputs = tf.convert_to_tensor(self.test_sentences)
            out = model(test_inputs)  # Build model with some sample inputs
            with TemporaryDirectory() as tempdir:
                save_path = Path(tempdir) / "saved.model"
                model.export(save_path)
                loaded_model = tf.saved_model.load(save_path)
            loaded_output = loaded_model.serve(test_inputs)
            # We may see small differences because the loaded model is compiled, so we need an epsilon for the test
            self.assertLessEqual(tf.reduce_max(tf.abs(out - loaded_output)), 1e-5)
//  [python]
def test_fine_tuning(self):
        model = CustomSignatureModel()
        model_no_signatures_path = os.path.join(
            self.get_temp_dir(), "model_no_signatures"
        )
        _ = model(tf.constant(0.0))

        tf.saved_model.save(model, model_no_signatures_path)
        restored_model = tf.compat.v1.saved_model.load(model_no_signatures_path)

        self.assertLen(list(restored_model.signatures.keys()), 0)
        self.assertEqual(restored_model(tf.constant(3.0)).numpy(), 3)
        restored_model.mutate(tf.constant(2.0))
        self.assertEqual(restored_model(tf.constant(3.0)).numpy(), 6)
        optimizer = optimizers.SGD(0.05)

        def train_step():
            with tf.GradientTape() as tape:
                loss = (10.0 - restored_model(tf.constant(2.0))) ** 2
            variables = tape.watched_variables()
            grads = tape.gradient(loss, variables)
            optimizer.apply_gradients(zip(grads, variables))
            return loss

        for _ in range(10):
            # "v" approaches 5, "loss" approaches 0
            loss = train_step()

        self.assertAllClose(loss, 0.0, rtol=1e-2, atol=1e-2)
        self.assertAllClose(restored_model.v.numpy(), 5.0, rtol=1e-2, atol=1e-2)
//  [python]
def test_fine_tuning(self):
        model = CustomSignatureModel()
        model_no_signatures_path = os.path.join(
            self.get_temp_dir(), "model_no_signatures"
        )
        _ = model(tf.constant(0.0))

        tf.saved_model.save(model, model_no_signatures_path)
        restored_model = tf.saved_model.load(model_no_signatures_path)

        self.assertLen(list(restored_model.signatures.keys()), 0)
        self.assertEqual(restored_model(tf.constant(3.0)).numpy(), 3)
        restored_model.mutate(tf.constant(2.0))
        self.assertEqual(restored_model(tf.constant(3.0)).numpy(), 6)
        optimizer = optimizers.SGD(0.05)

        def train_step():
            with tf.GradientTape() as tape:
                loss = (10.0 - restored_model(tf.constant(2.0))) ** 2
            variables = tape.watched_variables()
            grads = tape.gradient(loss, variables)
            optimizer.apply_gradients(zip(grads, variables))
            return loss

        for _ in range(10):
            # "v" approaches 5, "loss" approaches 0
            loss = train_step()

        self.assertAllClose(loss, 0.0, rtol=1e-2, atol=1e-2)
        self.assertAllClose(restored_model.v.numpy(), 5.0, rtol=1e-2, atol=1e-2)
//  [python]
def fit(self, training_data: PandasData, validation_data: PandasData) -> None:
        """Trains a tabnet model in a custom job.

        After the custom job successfully finishes, load the model and set it to
        self.model to enable prediction. If TensorFlow is not installed, the
        model will not be loaded.

        Training config can be overriden by setting the training config.

        Example Usage:
        `
        tabnet_trainer = TabNetTrainer(...)
        tabnet_trainer.fit.vertex.remote_config.staging_bucket = 'gs://...'
        tabnet_trainer.fit.vertex.remote_config.display_name = 'example'
        tabnet_trainer.fit(...)
        `

        PandasData refers to a pandas DataFrame. Each data frame should meet the
        following requirements:
            1. All entries should be numerical (no string, array or object).
            2. For categorical columns, the entries should be integers. In
            addition, the column type should be set to 'category'. Otherwise, it
            will be treated as numerical columns.
            3. The column names should be string.

        Args:
            training_data (pandas.DataFrame):
                Required. A pandas DataFrame for training.
            validation_data (pandas.DataFrame):
                Required. A pandas DataFrame for validation.
        """
        try:
            import tensorflow as tf

            self.model = tf.saved_model.load(self.output_model_path)
        except ImportError:
            _LOGGER.warning(
                "TensorFlow must be installed to load the trained model. The model is stored at %s",
                self.output_model_path,
            )
//  [python]
def fit(self, training_data: PandasData, validation_data: PandasData) -> None:
        """Trains a tabnet model in a custom job.

        After the custom job successfully finishes, load the model and set it to
        self.model to enable prediction. If TensorFlow is not installed, the
        model will not be loaded.

        Training config can be overriden by setting the training config.

        Example Usage:
        `
        tabnet_trainer = TabNetTrainer(...)
        tabnet_trainer.fit.vertex.remote_config.staging_bucket = 'gs://...'
        tabnet_trainer.fit.vertex.remote_config.display_name = 'example'
        tabnet_trainer.fit(...)
        `

        PandasData refers to a pandas DataFrame. Each data frame should meet the
        following requirements:
            1. All entries should be numerical (no string, array or object).
            2. For categorical columns, the entries should be integers. In
            addition, the column type should be set to 'category'. Otherwise, it
            will be treated as numerical columns.
            3. The column names should be string.

        Args:
            training_data (pandas.DataFrame):
                Required. A pandas DataFrame for training.
            validation_data (pandas.DataFrame):
                Required. A pandas DataFrame for validation.
        """
        try:
            import tensorflow.saved_model as tf_saved_model

            self.model = tf_saved_model.load(self.output_model_path)
        except ImportError:
            _LOGGER.warning(
                "TensorFlow must be installed to load the trained model. The model is stored at %s",
                self.output_model_path,
            )
//  [python]
def test_read_mirrored_variable(self, strategy):
    # synchronizaiton=ON_WRITE is the default variable created under
    # tf.distribute.Strategy.scope(). Most model parameters are this kind of
    # variable. Reading a synchronization=ON_WRITE simply reads the primary
    # component, so it works as intended.

    class Model(tf.Module):

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

      @tf.function(input_signature=[])
      def __call__(self):
        return self.v.read_value()

    export_dir = self.get_temp_dir()
    with strategy.scope():
      m = Model()
      m.v.assign(1.)
      tf.saved_model.save(m, export_dir)

    loaded = tf.compat.v1.saved_model.load(export_dir)
    self.assertEqual(self.evaluate(loaded()), 1.)
//  [python]
def test_read_mirrored_variable(self, strategy):
    # synchronizaiton=ON_WRITE is the default variable created under
    # tf.distribute.Strategy.scope(). Most model parameters are this kind of
    # variable. Reading a synchronization=ON_WRITE simply reads the primary
    # component, so it works as intended.

    class Model(tf.Module):

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

      @tf.function(input_signature=[])
      def __call__(self):
        return self.v.read_value()

    export_dir = self.get_temp_dir()
    with strategy.scope():
      m = Model()
      m.v.assign(1.)
      tf.saved_model.save(m, export_dir)

    loaded = tf.saved_model.load(export_dir)
    self.assertEqual(self.evaluate(loaded()), 1.)
//  [python]
def test_update_sync_on_read_variable(self, strategy):
    # It's rare to update aggregation=ON_READ variables in serving, but it's
    # possible that the SavedModel contains both serving and training graphs,
    # and the training may contain metrics layers.

    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=[])
      def update(self):
        self.v.assign_add(1.)
        return self.v.read_value()

    export_dir = self.get_temp_dir()
    with strategy.scope():
      m = Model()
      tf.saved_model.save(m, export_dir)

    loaded = tf.saved_model.load(export_dir)
    loaded.update()
    self.assertEqual(self.evaluate(loaded.v), 1.)
//  [python]
def test_read_sync_on_read_variable(self, strategy):
    # synchronizaiton=ON_READ variables are typically used in Keras metrics and
    # batch norm layers.

    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=[])
      def __call__(self):
        return self.v.read_value()

    export_dir = self.get_temp_dir()
    with strategy.scope():
      m = Model()
      # Note that each component is assigned with the value divided by the
      # number of replicas.
      m.v.assign(1.)
      self.assertAllEqual(
          self.evaluate(strategy.experimental_local_results(m.v)), [0.5, 0.5])
      tf.saved_model.save(m, export_dir)

    loaded = tf.compat.v1.saved_model.load(export_dir)
    # The variable already has the aggregated value.
    self.assertEqual(self.evaluate(loaded.v.read_value()), 1.)
    self.assertEqual(self.evaluate(loaded()), 1.)
//  [python]
def test_read_sync_on_read_variable(self, strategy):
    # synchronizaiton=ON_READ variables are typically used in Keras metrics and
    # batch norm layers.

    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=[])
      def __call__(self):
        return self.v.read_value()

    export_dir = self.get_temp_dir()
    with strategy.scope():
      m = Model()
      # Note that each component is assigned with the value divided by the
      # number of replicas.
      m.v.assign(1.)
      self.assertAllEqual(
          self.evaluate(strategy.experimental_local_results(m.v)), [0.5, 0.5])
      tf.saved_model.save(m, export_dir)

    loaded = tf.saved_model.load(export_dir)
    # The variable already has the aggregated value.
    self.assertEqual(self.evaluate(loaded.v.read_value()), 1.)
    self.assertEqual(self.evaluate(loaded()), 1.)
//  [python]
def test_update_mirrored_variable(self, strategy):
    # It's very rare to update aggregation=ON_WRITE variables in the forward
    # path, and this test case is mainly for completeness.

    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=[])
      def update(self):
        self.v.assign_add(1.)

    export_dir = self.get_temp_dir()
    with strategy.scope():
      m = Model()
      tf.saved_model.save(m, export_dir)

    loaded = tf.compat.v1.saved_model.load(export_dir)
    self.assertEqual(self.evaluate(loaded.v), 0.)
    loaded.update()
    self.assertEqual(self.evaluate(loaded.v), 1.)
//  [python]
def test_update_mirrored_variable(self, strategy):
    # It's very rare to update aggregation=ON_WRITE variables in the forward
    # path, and this test case is mainly for completeness.

    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=[])
      def update(self):
        self.v.assign_add(1.)

    export_dir = self.get_temp_dir()
    with strategy.scope():
      m = Model()
      tf.saved_model.save(m, export_dir)

    loaded = tf.saved_model.load(export_dir)
    self.assertEqual(self.evaluate(loaded.v), 0.)
    loaded.update()
    self.assertEqual(self.evaluate(loaded.v), 1.)