# Copyright 2022 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

import os
import tempfile

import numpy as np

from absl.testing import parameterized
from tensorflow.python.ipu import test_utils as tu
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.eager import def_function
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import function
from tensorflow.python.framework import test_util
from tensorflow.python.framework import tensor_spec
from tensorflow.python.ops.functional_ops import partitioned_call
from tensorflow.python.ipu import config
from tensorflow.python.ipu import ipu_strategy
from tensorflow.python.ipu import serving
from tensorflow.python.ipu.config import IPUConfig
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import variables
from tensorflow.python.platform import test
from tensorflow.python.saved_model import load
from tensorflow.python.saved_model import signature_constants
import popef


class TestServingExportBase(test_util.TensorFlowTestCase,
                            parameterized.TestCase):
  def setUp(self):
    super().setUp()
    self.use_ipus(1)

  def use_ipus(self, num):
    cfg = IPUConfig()
    cfg.device_connection.enable_remote_buffers = True
    cfg.device_connection.type = config.DeviceConnectionType.ON_DEMAND
    cfg.auto_select_ipus = num

    tu.add_hw_ci_connection_options(cfg)
    cfg.configure_ipu_system()

  def _load_and_run(self, path, inputs):
    imported = load.load(path)
    loaded = imported.signatures[
        signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY]

    input_names = [inp.name.split(':')[0] for inp in loaded.inputs]

    if isinstance(inputs, list):
      # Handle cases with input names generated by Keras.
      input_names = [name for name in input_names if "input_" in name]
      assert len(input_names) == len(inputs)
      input_names.sort(key=lambda name: int(name.split("_")[1]))
      inputs = {input_names[i]: inputs[i] for i in range(len(inputs))}
    elif not isinstance(inputs, dict):
      inputs = {input_names[0]: inputs}

    strategy = ipu_strategy.IPUStrategy()
    with strategy.scope():
      result = strategy.run(loaded, kwargs=inputs)

    return result


class TestServingExport(TestServingExportBase):
  @tu.test_uses_ipus(num_ipus=1, allow_ipu_model=False)
  @test_util.run_v2_only
  def test_export_programs_names(self):
    element_count = 3
    input_shape = (element_count,)
    input_signatures = (tensor_spec.TensorSpec(shape=input_shape,
                                               dtype=np.float16,
                                               name='in_x'),
                        tensor_spec.TensorSpec(shape=input_shape,
                                               dtype=np.float16,
                                               name='in_y'))
    var_value = np.float16(4.)
    w = variables.Variable(var_value)

    @def_function.function
    def my_net(x, y):
      return x * y + w

    with tempfile.TemporaryDirectory() as tmp_folder:
      iterations = 16
      serving.export_single_step(my_net,
                                 tmp_folder,
                                 iterations,
                                 input_signatures,
                                 output_names="out_z")

      assets_dir = os.path.join(tmp_folder, 'assets')
      popef_dir = os.path.join(assets_dir, os.listdir(assets_dir)[0])
      reader = popef.Reader()
      reader.parseFile(popef_dir)
      metadata = reader.metadata()[0]
      anchors = metadata.anchors()

      main_prog_idx = 1
      for idx in range(3):
        self.assertEqual(anchors[idx].programs(), [
            main_prog_idx,
        ])

      programs_map = metadata.programsMap()
      self.assertEqual(programs_map[0], 'load_program')
      self.assertEqual(programs_map[1], 'main_program')
      self.assertEqual(programs_map[2], 'save_program')

  @tu.test_uses_ipus(num_ipus=1, allow_ipu_model=False)
  @test_util.run_v2_only
  def test_export_simple_model_no_var_defunc(self):
    element_count = 4
    input_shape = (element_count,)
    predict_step_signature = (tensor_spec.TensorSpec(shape=input_shape,
                                                     dtype=np.float32,
                                                     name='x'),)

    @def_function.function
    def my_net(x):
      return x * x

    with tempfile.TemporaryDirectory() as tmp_folder:
      iterations = 16
      serving.export_single_step(my_net,
                                 tmp_folder,
                                 iterations,
                                 predict_step_signature,
                                 output_names="out0")

      input_data = np.arange(element_count, dtype=np.float32)

      # load and run
      result = self._load_and_run(tmp_folder, input_data)
      self.assertEqual(list(result["out0"]), list(np.square(input_data)))

  @tu.test_uses_ipus(num_ipus=1, allow_ipu_model=False)
  @test_util.run_v2_only
  def test_export_simple_model_no_var_no_defunc(self):
    element_count = 4
    input_shape = (element_count,)
    predict_step_signature = (tensor_spec.TensorSpec(shape=input_shape,
                                                     dtype=np.float32,
                                                     name='x'),)

    def my_net(x):
      return x * x

    with tempfile.TemporaryDirectory() as tmp_folder:
      iterations = 16
      serving.export_single_step(my_net,
                                 tmp_folder,
                                 iterations,
                                 predict_step_signature,
                                 output_names="out0")

      input_data = np.arange(element_count, dtype=np.float32)

      # load and run
      result = self._load_and_run(tmp_folder, input_data)
      self.assertEqual(list(result["out0"]), list(np.square(input_data)))

  @tu.test_uses_ipus(num_ipus=1, allow_ipu_model=False)
  @test_util.run_v2_only
  def test_export_simple_model_no_var_preprocessing(self):

    element_count = 15
    input_shape = (element_count,)
    predict_step_signature = (tensor_spec.TensorSpec(shape=input_shape,
                                                     dtype=np.float32,
                                                     name='x'),)

    preprocessing_step_signature = predict_step_signature

    @def_function.function
    def preprocessing_step(x):
      return x - 1.0

    @def_function.function
    def my_net(x):
      return x * x

    def expected_value(x):
      return list(np.square(np.subtract(input_data, 1.0)))

    with tempfile.TemporaryDirectory() as tmp_folder:
      iterations = 4
      serving.export_single_step(
          my_net,
          tmp_folder,
          iterations,
          predict_step_signature,
          preprocessing_step=preprocessing_step,
          preprocessing_step_signature=preprocessing_step_signature,
          output_names="out0")

      input_data = np.arange(element_count, dtype=np.float32)

      # load and run
      result = self._load_and_run(tmp_folder, input_data)
      self.assertEqual(list(result["out0"]), expected_value(input_data))

  @tu.test_uses_ipus(num_ipus=1, allow_ipu_model=False)
  @test_util.run_v2_only
  def test_export_simple_model_no_var_postprocessing(self):

    element_count = 15
    input_shape = (element_count,)
    predict_step_signature = (tensor_spec.TensorSpec(shape=input_shape,
                                                     dtype=np.float32,
                                                     name='x'),)

    postprocessing_step_signature = predict_step_signature

    @def_function.function
    def my_net(x):
      return x * x

    @def_function.function
    def postprocessing_step(x):
      return x - 1.0

    def expected_value(x):
      return list(np.subtract(np.square(input_data), 1.0))

    with tempfile.TemporaryDirectory() as tmp_folder:
      iterations = 4
      serving.export_single_step(
          my_net,
          tmp_folder,
          iterations,
          predict_step_signature,
          postprocessing_step=postprocessing_step,
          postprocessing_step_signature=postprocessing_step_signature,
          output_names="out0")

      input_data = np.arange(element_count, dtype=np.float32)

      # load and run
      result = self._load_and_run(tmp_folder, input_data)
      self.assertEqual(list(result["out0"]), expected_value(input_data))

  @tu.test_uses_ipus(num_ipus=1, allow_ipu_model=False)
  @test_util.run_v2_only
  def test_export_simple_model_no_var_pre_post_processing_transform_strings(
      self):

    element_count = 2
    input_shape = (element_count,)

    predict_step_signature = (tensor_spec.TensorSpec(shape=input_shape,
                                                     dtype=np.float32,
                                                     name='x'),)

    @def_function.function(input_signature=(tensor_spec.TensorSpec(
        shape=input_shape, dtype=dtypes.string, name="input string"),))
    def preprocessing_step(input_tensor):
      transform_fn = lambda input: constant_op.constant(
          1.0) if input == "graphcore" else constant_op.constant(2.0)

      return array_ops.stack(
          [transform_fn(elem) for elem in array_ops.unstack(input_tensor)])

    @def_function.function
    def my_net(x):
      return x * x

    @def_function.function(input_signature=(tensor_spec.TensorSpec(
        shape=(2,), dtype=dtypes.float32, name="input"),))
    def postprocessing_step(x):
      return x - 1.0

    input_data = constant_op.constant(["graphcore", "other"],
                                      shape=input_shape)

    with tempfile.TemporaryDirectory() as tmp_folder:
      iterations = 4
      serving.export_single_step(my_net,
                                 tmp_folder,
                                 iterations,
                                 predict_step_signature,
                                 preprocessing_step=preprocessing_step,
                                 postprocessing_step=postprocessing_step,
                                 output_names="out0")

      result = self._load_and_run(tmp_folder, input_data)
      self.assertEqual(list(result["out0"]), [0.0, 3.0])

  @tu.test_uses_ipus(num_ipus=1, allow_ipu_model=False)
  @test_util.run_v2_only
  def test_export_simple_model_with_variable(self):
    element_count = 3
    input_shape = (element_count,)
    input_tensor = array_ops.zeros(shape=input_shape, dtype=np.float16)

    dataset = dataset_ops.Dataset.from_tensors(input_tensor)

    var_value = np.float16(4.)
    w = variables.Variable(var_value)

    @def_function.function
    def my_net(x):
      return x * w

    with tempfile.TemporaryDirectory() as tmp_folder:
      iterations = 16
      serving.export_single_step(my_net,
                                 tmp_folder,
                                 iterations,
                                 input_dataset=dataset)

      input_data = np.arange(element_count, dtype=np.float16)
      result = self._load_and_run(tmp_folder, input_data)
      self.assertEqual(list(result["output_0"]), list(input_data * var_value))

  @tu.test_uses_ipus(num_ipus=1, allow_ipu_model=False)
  @test_util.run_v2_only
  def test_export_simple_model_with_variable_pre_post_processing(self):
    element_count = 17
    input_shape = (element_count,)
    input_tensor = array_ops.zeros(shape=input_shape, dtype=np.float16)
    predict_step_signature = tensor_spec.TensorSpec(shape=input_shape,
                                                    dtype=np.float16),

    dataset = dataset_ops.Dataset.from_tensors(input_tensor)

    preproc_value = np.float16(2.)
    var_value = np.float16(4.)
    w = variables.Variable(var_value)
    p = variables.Variable(preproc_value)

    @def_function.function
    def preprocessing_step(x):
      return x * p

    @def_function.function(input_signature=predict_step_signature)
    def my_net(x):
      return x * w

    @def_function.function(input_signature=predict_step_signature)
    def postprocessing_step(x):
      return x / p

    with tempfile.TemporaryDirectory() as tmp_folder:
      iterations = 16
      serving.export_single_step(my_net,
                                 tmp_folder,
                                 iterations,
                                 input_dataset=dataset,
                                 preprocessing_step=preprocessing_step,
                                 postprocessing_step=postprocessing_step)

      input_data = np.arange(element_count, dtype=np.float16)
      result = self._load_and_run(tmp_folder, input_data)
      self.assertEqual(list(result["output_0"]), list(input_data * var_value))

  @tu.test_uses_ipus(num_ipus=1, allow_ipu_model=False)
  @test_util.run_v2_only
  def test_export_simple_model_two_inputs(self):

    element_count = 3
    input_shape = (element_count,)
    predict_step_signature = (tensor_spec.TensorSpec(shape=input_shape,
                                                     dtype=np.float32),
                              tensor_spec.TensorSpec(shape=(),
                                                     dtype=np.float32))

    @def_function.function(input_signature=predict_step_signature)
    def my_net(x1, x2):
      return x1 * x2, x1 + x2

    with tempfile.TemporaryDirectory() as tmp_folder:
      iterations = 16
      serving.export_single_step(my_net,
                                 tmp_folder,
                                 iterations,
                                 output_names=["result0", "result1"])

      x1_data = np.arange(element_count, dtype=np.float32)
      x2_data = np.float32(3.0)
      result = self._load_and_run(tmp_folder, {'x1': x1_data, 'x2': x2_data})

      self.assertEqual(list(result["result0"]), list(x1_data * x2_data))
      self.assertEqual(list(result["result1"]), list(x1_data + x2_data))

  @tu.test_uses_ipus(num_ipus=1, allow_ipu_model=False)
  @test_util.run_v2_only
  def test_export_simple_model_two_inputs_pre_post_processing(self):

    element_count = 3
    input_shape = (element_count,)
    predict_step_signature = (tensor_spec.TensorSpec(shape=input_shape,
                                                     dtype=np.float32),
                              tensor_spec.TensorSpec(shape=(),
                                                     dtype=np.float32))

    postprocessing_step_signature = (tensor_spec.TensorSpec(shape=input_shape,
                                                            dtype=np.float32),
                                     tensor_spec.TensorSpec(shape=input_shape,
                                                            dtype=np.float32))

    @def_function.function(input_signature=predict_step_signature)
    def preprocessing_step(x1, x2):
      return x1 - 2, x2 + 5

    @def_function.function(input_signature=predict_step_signature)
    def my_net(x1, x2):
      return x1 * x2, x1 + x2

    @def_function.function(input_signature=postprocessing_step_signature)
    def postprocessing_step(x1, x2):
      return math_ops.reduce_sum(x1), math_ops.abs(x2)

    with tempfile.TemporaryDirectory() as tmp_folder:
      iterations = 16
      serving.export_single_step(my_net,
                                 tmp_folder,
                                 iterations,
                                 preprocessing_step=preprocessing_step,
                                 postprocessing_step=postprocessing_step,
                                 output_names=["result0", "result1"])

      x1_data = np.arange(element_count, dtype=np.float32)
      x2_data = np.float32(3.0)
      result = self._load_and_run(tmp_folder, {'x1': x1_data, 'x2': x2_data})

      x1_data -= 2
      x2_data += 5

      self.assertEqual(float(result["result0"]), np.sum(x1_data * x2_data))
      self.assertEqual(list(result["result1"]),
                       list(np.absolute(x1_data + x2_data)))

  @tu.test_uses_ipus(num_ipus=1, allow_ipu_model=False)
  @test_util.run_v2_only
  def test_purge_model_dir(self):
    input_signature = (tensor_spec.TensorSpec(shape=[], dtype=np.float32),)

    @def_function.function(input_signature=input_signature)
    def model_fn(x1):
      return x1 * 2.0

    with tempfile.TemporaryDirectory() as tmp_folder:
      serving.export_single_step(model_fn,
                                 tmp_folder,
                                 iterations=16,
                                 output_names=["result0"])

      # Export to the same dir with purge_export_dir=True and check if no exception is raised
      serving.export_single_step(model_fn,
                                 tmp_folder,
                                 iterations=16,
                                 output_names=["result0"],
                                 purge_export_dir=True)

  @tu.test_uses_ipus(num_ipus=1, allow_ipu_model=False)
  @test_util.run_v2_only
  def test_pipeline_purge_model_dir(self):
    input_signature = (tensor_spec.TensorSpec(shape=[], dtype=np.float32),)

    def stage1(x1):
      return x1 * 2.0

    def stage2(x1):
      return x1 * 3.0

    with tempfile.TemporaryDirectory() as tmp_folder:
      serving.export_pipeline([stage1, stage2],
                              tmp_folder,
                              iterations=16,
                              device_mapping=[0, 0],
                              predict_step_signature=input_signature,
                              output_names=["out0"])

      # Export to the same dir with purge_export_dir=True and check if no exception is raised
      serving.export_pipeline([stage1, stage2],
                              tmp_folder,
                              iterations=16,
                              device_mapping=[0, 0],
                              predict_step_signature=input_signature,
                              output_names=["out0"],
                              purge_export_dir=True)

  @tu.test_uses_ipus(num_ipus=1, allow_ipu_model=False)
  @test_util.run_v2_only
  def test_export_simple_model_run_locally(self):
    element_count = 3
    input_shape = (element_count,)

    inputs = (array_ops.zeros(input_shape, np.float32),
              array_ops.zeros(input_shape, np.float32))
    dataset = dataset_ops.Dataset.from_tensors(inputs)
    var_value = np.float32(4.)
    w = variables.Variable(var_value)

    @def_function.function
    def my_net(x1, x2):
      return (x1 + x2) * w

    with tempfile.TemporaryDirectory() as tmp_folder:
      iterations = 16
      runtime_func = serving.export_single_step(my_net,
                                                tmp_folder,
                                                iterations,
                                                input_dataset=dataset)

      x1_data = np.arange(element_count, dtype=np.float32)
      x2_data = x1_data * 2
      strategy = ipu_strategy.IPUStrategy()
      with strategy.scope():
        x1_data_tf = constant_op.constant(x1_data)
        x2_data_tf = constant_op.constant(x2_data)
        result = strategy.run(runtime_func, args=(x1_data_tf, x2_data_tf))
      result = result[0]
      ref_result = (x1_data + x2_data) * var_value
      self.assertEqual(list(result), list(ref_result))

  @tu.test_uses_ipus(num_ipus=1, allow_ipu_model=False)
  @test_util.run_v2_only
  def test_export_simple_model_run_locally_pre_post_processing(self):
    element_count = 3
    input_shape = (element_count,)

    predict_step_signature = (tensor_spec.TensorSpec(shape=input_shape,
                                                     dtype=np.float32),
                              tensor_spec.TensorSpec(shape=input_shape,
                                                     dtype=np.float32))
    postprocessing_signature = (tensor_spec.TensorSpec(shape=input_shape,
                                                       dtype=np.float32),)
    inputs = (array_ops.zeros(input_shape, np.float32),
              array_ops.zeros(input_shape, np.float32))
    dataset = dataset_ops.Dataset.from_tensors(inputs)
    var_value = np.float32(4.)
    w = variables.Variable(var_value)

    @def_function.function
    def preprocessing_step(x1, x2):
      return x1 - 2, x2 + 5

    @def_function.function(input_signature=predict_step_signature)
    def my_net(x1, x2):
      return (x1 + x2) * w

    @def_function.function(input_signature=postprocessing_signature)
    def postprocessing_step(x):
      return x * 10

    with tempfile.TemporaryDirectory() as tmp_folder:
      iterations = 16
      runtime_func = serving.export_single_step(
          my_net,
          tmp_folder,
          iterations,
          preprocessing_step=preprocessing_step,
          postprocessing_step=postprocessing_step,
          input_dataset=dataset)

      x1_data = np.arange(element_count, dtype=np.float32)
      x2_data = x1_data * 2
      strategy = ipu_strategy.IPUStrategy()
      with strategy.scope():
        x1_data_tf = constant_op.constant(x1_data)
        x2_data_tf = constant_op.constant(x2_data)
        result = strategy.run(runtime_func, args=(x1_data_tf, x2_data_tf))
      result = result[0]
      ref_result = ((x1_data - 2) + (x2_data + 5)) * var_value * 10
      self.assertEqual(list(result), list(ref_result))

  @tu.test_uses_ipus(num_ipus=1, allow_ipu_model=True)
  @test_util.run_v2_only
  def test_export_single_step_fails_for_non_empty_dir(self):
    predict_step_signature = (tensor_spec.TensorSpec(shape=(4,),
                                                     dtype=np.float32,
                                                     name='x'),)

    @def_function.function
    def my_net(x):
      return x * x

    with tempfile.TemporaryDirectory() as tmp_folder:
      open(os.path.join(tmp_folder, 'dummy_file'), 'w').close()
      with self.assertRaisesRegex(ValueError, "is not empty"):
        serving.export_single_step(my_net, tmp_folder, 16,
                                   predict_step_signature)

  @tu.test_uses_ipus(num_ipus=1, allow_ipu_model=False)
  @test_util.run_v2_only
  def test_export_pipeline(self):

    element_count = 4
    input_shape = (element_count,)
    predict_step_signature = (tensor_spec.TensorSpec(shape=input_shape,
                                                     dtype=np.float32),
                              tensor_spec.TensorSpec(shape=input_shape,
                                                     dtype=np.float32))

    var_value = np.float32(4.)
    w = variables.Variable(var_value)

    def stage1(x1, x2):
      return x1 * x2 + w

    def stage2(x):
      return x - 2 * w + 2

    with tempfile.TemporaryDirectory() as tmp_folder:
      serving.export_pipeline([stage1, stage2],
                              tmp_folder,
                              iterations=16,
                              device_mapping=[0, 0],
                              predict_step_signature=predict_step_signature)

      x1_data = np.arange(element_count, dtype=np.float32)
      x2_data = x1_data * 2

      result = self._load_and_run(tmp_folder, {'x1': x1_data, 'x2': x2_data})
      ref_result = x1_data * x2_data - var_value + 2
      self.assertEqual(list(result["output_0"]), list(ref_result))

  @tu.test_uses_ipus(num_ipus=1, allow_ipu_model=False)
  @test_util.run_v2_only
  def test_export_pipeline_pre_post_processing(self):

    element_count = 4
    input_shape = (element_count,)
    predict_step_signature = (tensor_spec.TensorSpec(shape=input_shape,
                                                     dtype=np.float32),
                              tensor_spec.TensorSpec(shape=input_shape,
                                                     dtype=np.float32))
    postprocessing_step_signature = (tensor_spec.TensorSpec(shape=input_shape,
                                                            dtype=np.float32),)
    var_value = np.float32(4.)
    w = variables.Variable(var_value)

    def preprocessing_step(x1, x2):
      return x1 * 10, x2 * x2

    def stage1(x1, x2):
      return x1 * x2 + w

    def stage2(x):
      return x - 2 * w + 2

    def postprocessing_step(x):
      return math_ops.reduce_sum(x)

    with tempfile.TemporaryDirectory() as tmp_folder:
      serving.export_pipeline(
          [stage1, stage2],
          tmp_folder,
          iterations=16,
          device_mapping=[0, 0],
          predict_step_signature=predict_step_signature,
          preprocessing_step=preprocessing_step,
          preprocessing_step_signature=predict_step_signature,
          postprocessing_step=postprocessing_step,
          postprocessing_step_signature=postprocessing_step_signature)

      x1_data = np.arange(element_count, dtype=np.float32)
      x2_data = x1_data * 2

      result = self._load_and_run(tmp_folder, {'x1': x1_data, 'x2': x2_data})
      ref_result = np.sum((x1_data * 10) * np.square(x2_data) - var_value + 2)
      self.assertEqual(float(result["output_0"]), float(ref_result))

  @tu.test_uses_ipus(num_ipus=1, allow_ipu_model=False)
  @test_util.run_v2_only
  def test_export_pipeline_tffunction_signature(self):

    element_count = 4
    input_shape = (element_count,)
    predict_step_signature = (tensor_spec.TensorSpec(shape=(),
                                                     dtype=np.float16),
                              tensor_spec.TensorSpec(shape=input_shape,
                                                     dtype=np.float16))

    @def_function.function(input_signature=predict_step_signature)
    def stage1(x1, x2):
      return x1 * x2

    def stage2(x):
      return x + 2, x * 3

    with tempfile.TemporaryDirectory() as tmp_folder:
      serving.export_pipeline([stage1, stage2],
                              tmp_folder,
                              iterations=16,
                              inputs=[np.float16(42.0)],
                              device_mapping=[0, 0],
                              output_names=["out0", "out1"])

      x2_data = np.arange(element_count, dtype=np.float16)
      result = self._load_and_run(tmp_folder, x2_data)
      ref_out0 = 42.0 * x2_data + 2
      ref_out1 = 42.0 * x2_data * 3
      self.assertEqual(list(result["out0"]), list(ref_out0))
      self.assertEqual(list(result["out1"]), list(ref_out1))

  @tu.test_uses_ipus(num_ipus=1, allow_ipu_model=False)
  @test_util.run_v2_only
  def test_export_pipeline_tffunction_signature_pre_post_processing(self):

    element_count = 4
    input_shape = (element_count,)
    predict_step_signature = (tensor_spec.TensorSpec(shape=(),
                                                     dtype=np.float16),
                              tensor_spec.TensorSpec(shape=input_shape,
                                                     dtype=np.float16))
    preprocessing_step_signature = tensor_spec.TensorSpec(shape=input_shape,
                                                          dtype=np.float16),

    postprocessing_step_signature = (tensor_spec.TensorSpec(shape=input_shape,
                                                            dtype=np.float16),
                                     tensor_spec.TensorSpec(shape=input_shape,
                                                            dtype=np.float16))

    @def_function.function(input_signature=preprocessing_step_signature)
    def preprocessing_step(x2):
      return x2 * x2

    @def_function.function(input_signature=predict_step_signature)
    def stage1(x1, x2):
      return x1 * x2

    def stage2(x):
      return x + 2, x * 3

    @def_function.function(input_signature=postprocessing_step_signature)
    def postprocessing_step(x1, x2):
      return math_ops.cast(x1,
                           dtypes.float32), math_ops.cast(x2, dtypes.float32)

    with tempfile.TemporaryDirectory() as tmp_folder:
      serving.export_pipeline([stage1, stage2],
                              tmp_folder,
                              iterations=16,
                              inputs=[np.float16(42.0)],
                              device_mapping=[0, 0],
                              preprocessing_step=preprocessing_step,
                              postprocessing_step=postprocessing_step,
                              output_names=["out0", "out1"])

      x2_data = np.arange(element_count, dtype=np.float16)
      result = self._load_and_run(tmp_folder, x2_data)

      x2_data = np.square(x2_data)
      ref_out0 = 42.0 * x2_data + 2
      ref_out1 = 42.0 * x2_data * 3
      self.assertEqual(list(result["out0"]), list(ref_out0.astype(np.float32)))
      self.assertEqual(list(result["out1"]), list(ref_out1.astype(np.float32)))

  @tu.test_uses_ipus(num_ipus=1, allow_ipu_model=False)
  @test_util.run_v2_only
  def test_export_pipeline_dataset_signature(self):

    element_count = 4
    input_shape = (element_count,)

    inputs = (array_ops.zeros(shape=input_shape, dtype=np.float32),
              array_ops.zeros(shape=(), dtype=np.float32))
    dataset = dataset_ops.Dataset.from_tensors(inputs)

    def stage1(x1, x2, x3):
      return x1 * x2 + x3

    def stage2(x):
      return x + 2

    with tempfile.TemporaryDirectory() as tmp_folder:
      serving.export_pipeline([stage1, stage2],
                              tmp_folder,
                              iterations=16,
                              inputs=[42.0],
                              device_mapping=[0, 0],
                              input_dataset=dataset)

      x2_data = np.arange(element_count, dtype=np.float32)
      x3_data = np.float32(5.0)
      result = self._load_and_run(tmp_folder, {'x2': x2_data, 'x3': x3_data})
      ref_result = 42.0 * x2_data + x3_data + 2
      self.assertEqual(list(result["output_0"]), list(ref_result))

  @tu.test_uses_ipus(num_ipus=1, allow_ipu_model=False)
  @test_util.run_v2_only
  def test_export_pipeline_dataset_signature_pre_post_processing(self):

    element_count = 4
    input_shape = (element_count,)

    inputs = (array_ops.zeros(shape=input_shape, dtype=np.float32),
              array_ops.zeros(shape=(), dtype=np.float32))
    dataset = dataset_ops.Dataset.from_tensors(inputs)
    predict_step_signature = (tensor_spec.TensorSpec(shape=(),
                                                     dtype=np.float32),
                              tensor_spec.TensorSpec(shape=input_shape,
                                                     dtype=np.float32),
                              tensor_spec.TensorSpec(shape=(),
                                                     dtype=np.float32))
    postprocessing_step_signature = (tensor_spec.TensorSpec(shape=input_shape,
                                                            dtype=np.float32),)

    def preprocessing_step(x2, x3):
      return x2 * x3, x3 + 5

    @def_function.function(input_signature=predict_step_signature)
    def stage1(x1, x2, x3):
      return x1 * x2 + x3

    def stage2(x):
      return x + 2

    def postprocessing_step(x):
      return x - 20

    with tempfile.TemporaryDirectory() as tmp_folder:
      serving.export_pipeline(
          [stage1, stage2],
          tmp_folder,
          iterations=16,
          inputs=[42.0],
          device_mapping=[0, 0],
          input_dataset=dataset,
          preprocessing_step=preprocessing_step,
          postprocessing_step=postprocessing_step,
          postprocessing_step_signature=postprocessing_step_signature)

      x2_data = np.arange(element_count, dtype=np.float32)
      x3_data = np.float32(5.0)
      result = self._load_and_run(tmp_folder, {'x2': x2_data, 'x3': x3_data})
      ref_result = (42.0 * (x2_data * x3_data) + (x3_data + 5) + 2) - 20
      self.assertEqual(list(result["output_0"]), list(ref_result))

  @tu.test_uses_ipus(num_ipus=1, allow_ipu_model=True)
  @test_util.run_v2_only
  def test_export_pipeline_fails_for_non_empty_dir(self):
    predict_step_signature = (tensor_spec.TensorSpec(shape=(4,),
                                                     dtype=np.float32))

    def stage(x):
      return x + 2

    with tempfile.TemporaryDirectory() as tmp_folder:
      open(os.path.join(tmp_folder, 'dummy_file'), 'w').close()
      with self.assertRaisesRegex(ValueError, "is not empty"):
        serving.export_pipeline([stage, stage],
                                tmp_folder,
                                iterations=16,
                                device_mapping=[0, 0],
                                predict_step_signature=predict_step_signature)

  @tu.test_uses_ipus(num_ipus=1, allow_ipu_model=False)
  @test_util.run_v2_only
  def test_export_simple_model_with_variables_stateful_partition_call(self):
    element_count = 3
    input_shape = (element_count,)
    input_tensor = array_ops.zeros(shape=input_shape, dtype=np.float16)

    predict_step_signature = (tensor_spec.TensorSpec(shape=input_shape,
                                                     dtype=dtypes.float32),)

    var_value = np.float32(10.)
    w = variables.Variable(var_value)

    z_value = np.float32(15.)
    z = variables.Variable(z_value)

    @function.Defun(dtypes.float32, dtypes.float32)
    def partition_call_body(constant, var):
      return (constant + constant) * var

    @def_function.function(
        input_signature=(tensor_spec.TensorSpec(shape=input_shape,
                                                dtype=dtypes.float32),
                         tensor_spec.TensorSpec(shape=input_shape,
                                                dtype=dtypes.float32)))
    def predict_step(x, y):
      call_result = partitioned_call(args=[x, w], f=partition_call_body)

      predict_step_result = call_result + y + z
      return array_ops.reshape(predict_step_result, input_shape)

    x = np.arange(element_count, dtype=np.float32)
    y = np.arange(element_count, dtype=np.float32)
    expected_result = list(((x + x) * var_value) + y + z_value)

    with tempfile.TemporaryDirectory() as tmp_folder:
      iterations = 16
      serving.export_single_step(predict_step, tmp_folder, iterations)

      result = list(
          self._load_and_run(tmp_folder, {
              'x': x,
              'y': y
          })["output_0"].numpy())
      self.assertAllClose(result, expected_result)


if __name__ == "__main__":
  test.main()
