From 8f219b95bbce52a5cc447817af2f123732ec8c08 Mon Sep 17 00:00:00 2001
From: Alexander Pivovarov <pivovaa@amazon.com>
Date: Mon, 10 Jun 2019 15:29:01 -0700
Subject: [PATCH] Add PAD operator to relay tflite frontend (#3310)

---
 python/tvm/relay/frontend/tflite.py          | 26 ++++++++++++++++++++++++
 tests/python/frontend/tflite/test_forward.py | 30 ++++++++++++++++++++++++++++
 2 files changed, 56 insertions(+)

diff --git a/python/tvm/relay/frontend/tflite.py b/python/tvm/relay/frontend/tflite.py
index 9c8f50f..3a13473 100644
--- a/python/tvm/relay/frontend/tflite.py
+++ b/python/tvm/relay/frontend/tflite.py
@@ -66,6 +66,7 @@ class OperatorConverter(object):
             'ADD': self.convert_add,
             'MUL': self.convert_mul,
             'FULLY_CONNECTED': self.convert_fully_connected,
+            'PAD': self.convert_pad,
         }
 
     def check_unsupported_ops(self):
@@ -596,6 +597,31 @@ class OperatorConverter(object):
 
         return out
 
+    def convert_pad(self, op):
+        """Convert TFLite PAD"""
+        try:
+            from tflite.Operator import Operator
+        except ImportError:
+            raise ImportError("The tflite package must be installed")
+
+        assert isinstance(op, Operator)
+        input_tensors = self.get_input_tensors(op)
+        assert len(input_tensors) == 2, "input tensors length should be 2"
+
+        # TFLite only support CONSTANT mode and does not support constant_values parameter.
+        # tensor
+        input_tensor = input_tensors[0]
+        in_expr = self.get_expr(input_tensor.tensor_idx)
+
+        # paddings
+        pad_list = self.get_tensor_value(input_tensors[1])
+        # convert list of lists to tuple of tuples
+        paddings = tuple(tuple(l) for l in pad_list)
+
+        # Use default pad_value 0 because TFLite does not support constant_values parameter
+        out = _op.nn.pad(in_expr, paddings)
+        return out
+
     def get_expr(self, input_tensor_idx):
         return self.exp_tab.get_expr(get_tensor_name(self.subgraph, input_tensor_idx))
 
diff --git a/tests/python/frontend/tflite/test_forward.py b/tests/python/frontend/tflite/test_forward.py
index 677fbb8..7da2b85 100644
--- a/tests/python/frontend/tflite/test_forward.py
+++ b/tests/python/frontend/tflite/test_forward.py
@@ -394,6 +394,35 @@ def test_forward_squeeze():
     _test_squeeze(np.arange(6).reshape((1, 2, 1, 3)), [0, 2])
     _test_squeeze(np.arange(6).reshape((2, 1, 3, 1)), [1, 3])
 
+
+#######################################################################
+# Pad
+# ---
+
+def _test_pad(data):
+    """ One iteration of PAD """
+
+    assert len(data) == 2
+
+    # Test with tensor and constant
+    with tf.Graph().as_default():
+        in_data = [array_ops.placeholder(shape=data[0].shape, dtype=data[0].dtype, name='in')]
+        out = array_ops.pad(in_data[0], ops.convert_to_tensor(data[1], dtype=data[1].dtype))
+        compare_tflite_with_tvm([data[0]], ['in:0'], in_data, [out])
+
+
+def test_forward_pad():
+    """ Pad """
+    _test_pad([np.arange(1.0, 7.0, dtype=np.float32).reshape((2, 1, 1, 3)),
+               np.array([[1, 1], [2, 2], [1, 1], [2, 2]], dtype=np.int32)])
+    _test_pad([np.arange(1.0, 7.0, dtype=np.float32).reshape((2, 1, 3)),
+               np.array([[2, 2], [1, 1], [1, 1]], dtype=np.int32)])
+    _test_pad([np.arange(1.0, 7.0, dtype=np.float32).reshape((2, 3)),
+               np.array([[1, 1], [2, 2]], dtype=np.int32)])
+    _test_pad([np.arange(1.0, 4.0, dtype=np.float32).reshape((1, 3)),
+               np.array([[1, 1], [2, 2]], dtype=np.int32)])
+
+
 #######################################################################
 # Softmax
 # -------
@@ -528,6 +557,7 @@ def test_forward_inception_v4_net():
 if __name__ == '__main__':
     # Transforms
     test_forward_concatenation()
+    test_forward_pad()
     test_forward_reshape()
     test_forward_squeeze()
 
-- 
2.7.4