[Relay][Frontend] Support TF Gather

Author: yongwww

python/tvm/relay/frontend/tensorflow.py

@@ -676,6 +676,19 @@ def _impl(inputs, attr, params):
                                  'Taxis', '_class'])(new_input, attr)
     return _impl
 
+def _gather_v1():
+    def _impl(inputs, attr, params):
+        #axis = params[inputs.pop(2).name_hint].asnumpy()[0]
+        axis = 0
+        new_input = []
+        new_input.append(inputs.pop(0))
+        new_input.append(inputs.pop(0))
+        return AttrCvt(op_name="take",
+                       extras={'axis': tvm.const(axis, 'int32')},
+                       ignores=['Tindices', 'Tparams', 'validate_indices', \
+                                'Taxis', '_class'])(new_input, attr)
+     return _impl
+
 def _infer_out_shapes(inputs, params):
     """A method to get the output shape of an intermediate node in the relay graph."""
     out_type = ir_pass.infer_type(inputs)
@@ -1003,6 +1016,7 @@ def _impl(inputs, attr, params):
     'Sigmoid'                           : AttrCvt('sigmoid'),
     'Fill'                              : _fill(),
     'GatherV2'                          : _gather_v2(),
+    'Gather'                            : _gather_v1(),
     'StridedSlice'                      : _stridedSlice(),
     'LRN'                               : _lrn(),
     'Pad'                               : _pad('Pad'),

tests/python/frontend/tensorflow/test_forward.py

@@ -473,11 +473,11 @@ def test_forward_stridedslice():
 
 
 #######################################################################
-# Gather
+# Gather, GatherV2
 # ------
 
 def _test_gather(ip_shape, indice_shape, indice_value, axis, dtype):
-    """ One iteration of a Gather """
+    """ One iteration of a GatherV2 """
 
     tf.reset_default_graph()
     in_data = tf.placeholder(dtype, ip_shape, name="in_data")
@@ -497,7 +497,7 @@ def _fill_indices(indice_value):
     compare_tf_with_tvm([np_data, np_indices], ['in_data:0', 'indices:0'], 'GatherV2:0')
 
 def test_forward_gather():
-    '''test gather layer'''
+    '''test GatherV2 layer'''
     _test_gather((4,), (1,), 1, 0, 'int32')
     _test_gather((4,), (1,), 1, 0, 'float32')
     _test_gather((1,4), (1,), [0], 0, 'int32')
@@ -509,6 +509,42 @@ def test_forward_gather():
     _test_gather((3,3,3), (1,1,2), [[[1,0]]], 2, 'int32')
     _test_gather((4,3,5,6), (1,4), [[2,1,0,0]], 0, 'float32')
 
+
+def _test_gather_v1(ip_shape, indice_shape, indice_value, dtype):
+    """ One iteration of a Gather"""
+    tf.reset_default_graph()
+    in_data = tf.placeholder(dtype, ip_shape, name="in_data")
+    indices = tf.placeholder("int32", indice_shape, name="indices")
+    tf.gather(in_data, indices)
+    np_data = np.random.uniform(size=ip_shape).astype(dtype)
+
+    def _fill_indices(indice_value):
+        indices = np.array(ip_shape, dtype=dtype)
+        if isinstance(indice_value, int):
+            indices = np.array([indice_value], dtype='int32')
+        else:
+            indices = np.asarray(indice_value, dtype='int32')
+        return indices
+    np_indices = _fill_indices(indice_value)
+
+    compare_tf_with_tvm([np_data, np_indices], ['in_data:0', 'indices:0'], 'Gather:0')
+
+def test_forward_gather_v1():
+    '''test gather layer'''
+    #_test_gather((2,3), (1,), 1, 0, 'int32')
+    _test_gather_v1((4,), (1, 2, 2), [[[1, 0], [0, 1]]], 'float32')
+    _test_gather_v1((4,), (1,), 1, 'int32')
+    _test_gather_v1((4,), (1,), 1, 'float32')
+    _test_gather_v1((1,4), (1,), [0], 'int32')
+    _test_gather_v1((4,), (1,2,2), [[[1,0],[0,1]]], 'float32')
+    _test_gather_v1((2,2), (1,2,2), [[[1,0],[0,1]]], 'int32')
+    _test_gather_v1((2,2), (1,2,2), [[[1,0],[0,1]]], 'int32')
+    _test_gather_v1((2,2), (1,2,2), [[[1,0],[0,1]]], 'float32')
+    _test_gather_v1((3,3,3), (1,1,2), [[[1,0]]], 'int32')
+    _test_gather_v1((3,3,3), (1,1,2), [[[1,0]]], 'int32')
+    _test_gather_v1((4,3,5,6), (1,4), [[2,1,0,0]], 'float32')
+
+
 #######################################################################
 # Split
 # -----
@@ -1213,6 +1249,8 @@ def test_forward_rel_ops():
     test_forward_crop()
     test_forward_pad()
     test_forward_gather()
+    # Gather was used in tf 1.6 and before
+    # test_forward_gather_v1()
     test_forward_stridedslice()
     test_forward_split()
     test_forward_unstack()