[PaddlePaddle Hackathon 4][Frontend][Paddle]add thresholded_relu/index_select/eye/linspace/take_alone_axis/dist for paddle frontend

python/tvm/relay/frontend/paddlepaddle.py

@@ -400,6 +400,30 @@ def convert_conv2d_transpose(g, op, block):
     g.add_node(op.output("Output")[0], out)
 
 
+def convert_dist(g, op, block):
+    """Operator converter for dist."""
+
+    x = g.get_node(op.input("X")[0])
+    y = g.get_node(op.input("Y")[0])
+    z = _op.abs(_op.subtract(x, y))
+    dtype = infer_type(x).checked_type.dtype
+    p = op.attr("p")
+    if p == np.inf:
+        out = _op.reduce.max(_op.abs(z))
+    elif p == np.NINF:
+        out = _op.reduce.min(_op.abs(z))
+    elif p == 0.0:
+        out = _op.reduce.sum(_op.sign(_op.abs(z)))
+    else:
+        inv_p = _expr.const(1.0 / p, dtype=dtype)
+        p = _expr.const(p, dtype=dtype)
+        power_z = _op.power(z, p)
+        sum_pow = _op.reduce.sum(power_z)
+        out = _op.power(sum_pow, inv_p)
+    out = _op.full(out, shape=(1))
+    g.add_node(op.output("Out")[0], out)
+
+
 def convert_cumsum(g, op, block):
     """Operator converter for cumsum."""
 
@@ -475,6 +499,50 @@ def convert_elementwise_op(g, op, block):
     g.add_node(op.output("Out")[0], out)
 
 
+def convert_linspace(g, op, block):
+    """Operator converter for linspace."""
+
+    start = g.get_node(op.input("Start")[0])
+    stop = g.get_node(op.input("Stop")[0])
+    num = g.get_node(op.input("Num")[0])
+    dtype = _convert_dtype_value(op.attr("dtype"))
+    start, infered = try_infer_value(start, parameters=g.get_params())
+    if infered:
+        start = start.tolist()[0]
+    else:
+        msg = 'Value {} in attribute "start" of operator Linspace is not "valid."'
+        raise tvm.error.OpAttributeInvalid(msg.format(start))
+
+    stop, infered = try_infer_value(stop, parameters=g.get_params())
+    if infered:
+        stop = stop.tolist()[0]
+    else:
+        msg = 'Value {} in attribute "stop" of operator Linspace is not "valid."'
+        raise tvm.error.OpAttributeInvalid(msg.format(stop))
+
+    num, infered = try_infer_value(num, parameters=g.get_params())
+    if infered:
+        num = num.tolist()[0]
+    else:
+        msg = 'Value {} in attribute "num" of operator Linspace is not "valid."'
+        raise tvm.error.OpAttributeInvalid(msg.format(num))
+
+    if num == 1:
+        out = _op.full(_expr.const(start, dtype), shape=(1))
+    else:
+        if dtype in ["int32", "int64"]:
+            start = int(start)
+            stop = int(stop)
+        step = (stop - start) / (num - 1)
+        stop = stop + step
+        start = _expr.const(start, "float32")
+        stop = _expr.const(stop, "float32")
+        step = _expr.const(step, "float32")
+        out = _op.transform.arange(start=start, stop=stop, step=step, dtype="float32")
+        out = _op.cast(out, dtype)
+    g.add_node(op.output("Out")[0], out)
+
+
 def convert_elu(g, op, block):
     """Operator converter for elu."""
 
@@ -514,6 +582,27 @@ def convert_expand_as(g, op, block):
     g.add_node(op.output("Out")[0], out)
 
 
+def convert_eye(g, op, block):
+    """Operator converter for eye."""
+
+    num_rows = op.attr("num_rows")
+    num_columns = op.attr("num_columns")
+    one_nums = min(num_rows, num_columns)
+    dtype = op.attr("dtype")
+    dtype = _convert_dtype_value(dtype)
+
+    zeros = _op.zeros((num_rows, num_columns), dtype)
+    if one_nums == 0:
+        out = zeros
+    else:
+        ones = _op.ones(one_nums, dtype)
+        indices = _op.arange(
+            _expr.const(0, dtype="int32"), _expr.const(one_nums, dtype="int32"), dtype="int32"
+        )
+        out = _op.scatter_nd(zeros, _op.stack([indices, indices], axis=0), ones, "update")
+    g.add_node(op.output("Out")[0], out)
+
+
 def convert_feed(g, op, block):
     """Converter for model input node."""
 
@@ -830,6 +919,16 @@ def get_interpolate_mode(op):
     g.add_node(op.output("Out")[0], out)
 
 
+def convert_index_select(g, op, block):
+    """Operator converter for index_select."""
+
+    x = g.get_node(op.input("X")[0])
+    index = g.get_node(op.input("Index")[0])
+    axis = op.attr("dim")
+    out = _op.transform.take(x, index, axis, mode="wrap")
+    g.add_node(op.output("Out")[0], out)
+
+
 def convert_instance_norm(g, op, block):
     """Operator converter for instance_norm."""
 
@@ -2071,6 +2170,28 @@ def convert_swish(g, op, block):
     g.add_node(op.output("Out")[0], out)
 
 
+def convert_take_along_axis(g, op, block):
+    """Operator converter for take_along_axis."""
+
+    x = g.get_node(op.input("Input")[0])
+    index = g.get_node(op.input("Index")[0])
+    axis = op.attr("Axis")
+    out = _op.gather(x, axis, index)
+    g.add_node(op.output("Result")[0], out)
+
+
+def convert_thresholded_relu(g, op, block):
+    """Operator converter for thresholded_relu."""
+
+    x = g.get_node(op.input("X")[0])
+    dtype = infer_type(x).checked_type.dtype
+    threshold = op.attr("threshold")
+    threshold = _expr.const(threshold, dtype)
+    zero = _expr.const(0, dtype=dtype)
+    out = tvm.relay.where(x > threshold, x, zero)
+    g.add_node(op.output("Out")[0], out)
+
+
 def convert_tile(g, op, block):
     """Operator converter for tile."""
 
@@ -2208,6 +2329,7 @@ def convert_where_index(g, op, block):
     "cumsum": convert_cumsum,
     "depthwise_conv2d": convert_conv2d,
     "depthwise_conv2d_transpose": convert_conv2d_transpose,
+    "dist": convert_dist,
     "dot": convert_dot,
     "dropout": convert_dropout,
     "elementwise_add": convert_elementwise_op,
@@ -2226,6 +2348,7 @@ def convert_where_index(g, op, block):
     "exp": convert_unary_op,
     "expand_v2": convert_expand,
     "expand_as_v2": convert_expand_as,
+    "eye": convert_eye,
     "feed": convert_feed,
     "fill_any_like": convert_fill_any_like,
     "fill_constant": convert_fill_constant,
@@ -2242,6 +2365,7 @@ def convert_where_index(g, op, block):
     "hard_shrink": convert_hard_shrink,
     "hard_sigmoid": convert_hard_sigmoid,
     "hard_swish": convert_hard_swish,
+    "index_select": convert_index_select,
     "instance_norm": convert_instance_norm,
     "isfinite_v2": convert_unary_op,
     "isinf_v2": convert_unary_op,
@@ -2250,6 +2374,7 @@ def convert_where_index(g, op, block):
     "leaky_relu": convert_leaky_relu,
     "less_equal": convert_elementwise_op,
     "less_than": convert_elementwise_op,
+    "linspace": convert_linspace,
     "log": convert_unary_op,
     "log2": convert_unary_op,
     "log10": convert_unary_op,
@@ -2317,8 +2442,10 @@ def convert_where_index(g, op, block):
     "square": convert_square,
     "squeeze2": convert_squeeze,
     "swish": convert_swish,
+    "take_along_axis": convert_take_along_axis,
     "tan": convert_unary_op,
     "tanh": convert_unary_op,
+    "thresholded_relu": convert_thresholded_relu,
     "tile": convert_tile,
     "top_k_v2": convert_topk,
     "transpose2": convert_transpose,

tests/python/frontend/paddlepaddle/test_forward.py

@@ -1969,5 +1969,142 @@ def forward(self, inputs):
         verify_model(Mish(), input_data=input_data)
 
 
+@tvm.testing.uses_gpu
+def test_forward_thresholded_relu():
+    class ThresholdedRelu(nn.Layer):
+        @paddle.jit.to_static
+        def forward(self, inputs):
+            return nn.functional.thresholded_relu(inputs)
+
+    input_shapes = [[10], [2, 3], [5, 10, 11], [3, 4, 5, 6]]
+    for input_shape in input_shapes:
+        input_data = paddle.randn(shape=input_shape, dtype="float32")
+        verify_model(ThresholdedRelu(), input_data=input_data)
+
+
+@tvm.testing.uses_gpu
+def test_forward_index_select():
+    class IndexSelect1(nn.Layer):
+        @paddle.jit.to_static
+        def forward(self, x, index):
+            return paddle.index_select(x, index, axis=0)
+
+    class IndexSelect2(nn.Layer):
+        @paddle.jit.to_static
+        def forward(self, x, index):
+            return paddle.index_select(x, index, axis=-1)
+
+    input_shapes = [[10], [2, 3], [5, 10, 11], [3, 4, 5, 6]]
+    for input_shape in input_shapes:
+        input_data = paddle.randn(shape=input_shape, dtype="float32")
+        index = paddle.to_tensor([0, 1, 1], dtype="int32")
+        verify_model(IndexSelect1(), input_data=[input_data, index])
+        verify_model(IndexSelect2(), input_data=[input_data, index])
+
+
+@tvm.testing.uses_gpu
+def test_forward_eye():
+    class Eye1(nn.Layer):
+        @paddle.jit.to_static
+        def forward(self, inputs):
+            return paddle.eye(3, 5, dtype="int32"), paddle.eye(3, 5, dtype="float32"), inputs
+
+    class Eye2(nn.Layer):
+        @paddle.jit.to_static
+        def forward(self, inputs):
+            return paddle.eye(5, 3, dtype="int64"), paddle.eye(5, 3, dtype="float64"), inputs
+
+    class Eye3(nn.Layer):
+        @paddle.jit.to_static
+        def forward(self, inputs):
+            return paddle.eye(0, 3, dtype="int64"), paddle.eye(0, 0, dtype="float64"), inputs
+
+    x = paddle.to_tensor([1], dtype="float32")
+    verify_model(Eye1(), input_data=[x])
+    verify_model(Eye2(), input_data=[x])
+    verify_model(Eye3(), input_data=[x])
+
+
+@tvm.testing.uses_gpu
+def test_forward_linspace():
+    class Linspace1(nn.Layer):
+        @paddle.jit.to_static
+        def forward(self, inputs):
+            out1 = paddle.linspace(0.5, 7, 1, "int32")
+            out2 = paddle.linspace(1.3, 7.1, 5, "float32")
+            return out1, out2, inputs
+
+    class Linspace2(nn.Layer):
+        @paddle.jit.to_static
+        def forward(self, inputs):
+            start = paddle.to_tensor([2.5])
+            stop = paddle.to_tensor([3.6])
+            num = paddle.to_tensor([3])
+            start = paddle.cast(start, "float32")
+            stop = paddle.cast(stop, "float32")
+            num = paddle.cast(num, "int32")
+            out1 = paddle.linspace(start, stop, num, "int32")
+            out2 = paddle.linspace(start, stop, num, "float32")
+            return out1, out2, inputs
+
+    x = paddle.to_tensor([1], dtype="float32")
+    verify_model(Linspace1(), input_data=[x])
+    verify_model(Linspace2(), input_data=[x])
+
+
+@tvm.testing.uses_gpu
+def test_forward_take_alone_axis():
+    class TakeAloneAxis1(nn.Layer):
+        @paddle.jit.to_static
+        def forward(self, inputs):
+            index = paddle.to_tensor([[0, 1], [1, 2], [1, 0]])
+            return paddle.take_along_axis(inputs, index, axis=-1)
+
+    class TakeAloneAxis2(nn.Layer):
+        @paddle.jit.to_static
+        def forward(self, inputs):
+            index = paddle.to_tensor([[[0], [0], [1]]])
+            return paddle.take_along_axis(inputs, index, axis=1)
+
+    class TakeAloneAxis3(nn.Layer):
+        @paddle.jit.to_static
+        def forward(self, inputs):
+            index = paddle.to_tensor([[[0], [0], [1]]])
+            return paddle.take_along_axis(inputs, index, axis=-1)
+
+    if paddle.version.full_version >= "2.4.2":
+        x = paddle.to_tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
+        verify_model(TakeAloneAxis1(), input_data=x)
+
+        y = paddle.to_tensor(
+            [[[1, 2], [2, 3], [3, 4]], [[4, 5], [5, 6], [6, 7]], [[7, 8], [8, 9], [9, 10]]],
+            dtype="float32",
+        )
+        verify_model(TakeAloneAxis2(), input_data=y)
+        verify_model(TakeAloneAxis3(), input_data=y)
+
+
+@tvm.testing.uses_gpu
+def test_forward_dist():
+    class Dist(nn.Layer):
+        @paddle.jit.to_static
+        def forward(self, x, y):
+            l0_norm = paddle.dist(x, y, 0)
+            l2_norm = paddle.dist(x, y, 2)
+            float_norm = paddle.dist(x, y, 1.3)
+            inf_norm = paddle.dist(x, y, float("inf"))
+            ninf_norm = paddle.dist(x, y, float("-inf"))
+            return l0_norm, l2_norm, float_norm, inf_norm, ninf_norm
+
+    x = paddle.to_tensor([[3, 3], [3, 3]], dtype="float32")
+    y = paddle.to_tensor([[1, 2], [3, 4]], dtype="float32")
+    w = paddle.to_tensor([[1, 2]], dtype="float32")
+    v = paddle.to_tensor([[2.1]], dtype="float32")
+    verify_model(Dist(), input_data=[x, y])
+    verify_model(Dist(), input_data=[x, w])
+    verify_model(Dist(), input_data=[w, v])
+    verify_model(Dist(), input_data=[y, v])
+
+
 if __name__ == "__main__":
     tvm.testing.main()