[Unity][Frontend] Add Sigmoid and Square Op

python/tvm/relax/frontend/nn/op.py

@@ -843,6 +843,31 @@ def gelu(x: Tensor, approximate: Optional[str] = None, name: str = "gelu") -> Te
     return wrap_nested(gelu_out, name)
 
 
+def sigmoid(x: Tensor, name: str = "sigmoid") -> Tensor:
+    r"""Computes sigmoid.
+
+    .. math:: \text{sigmoid}(x) = \frac{1}{1 + \exp(-x)}
+
+    Parameters
+    ----------
+    data: Tensor
+        The input data to the operator.
+
+    name : str
+        Name hint.
+
+    Returns
+    -------
+    result : Tensor
+        The computed result.
+
+    Note
+    ----
+    The input tensor is required to have float dtype
+    """
+    return wrap_nested(_op.sigmoid(x._expr), name)
+
+
 def softmax(x: Tensor, axis: int = -1, name: str = "softmax") -> Tensor:
     r"""Computes softmax.
 
@@ -1268,6 +1293,25 @@ def pad(
     return wrap_nested(_op.nn.pad(x._expr, pad_width=pad, pad_value=value, pad_mode=mode), name)
 
 
+def square(x: Tensor, name: str = "square") -> Tensor:
+    """Computes the element-wise square of the input tensor.
+
+    Parameters
+    ----------
+    x : Tensor
+        The input tensor.
+
+    name : str
+        Name hint.
+
+    Returns
+    -------
+    result : Tensor
+        The computed result.
+    """
+    return wrap_nested(_op.square(x._expr), name)
+
+
 def get_timestep_embedding(
     x: Tensor,
     embedding_dim: int,

tests/python/relax/test_frontend_nn_op.py

@@ -26,6 +26,31 @@
 # mypy: disable-error-code="attr-defined,valid-type,name-defined"
 
 
+def test_unary():
+    class Model(Module):
+        def test(self, x: Tensor):
+            z0 = op.square(x)
+            return (x,)
+
+    # fmt: off
+    @R.function
+    def test(x: R.Tensor((1, 10), dtype="float32"), _io: R.Object):
+        R.func_attr({"num_input": 2})
+        with R.dataflow():
+            square: R.Tensor((1, 10), dtype="float32") = R.square(x)
+            gv1 = (x,), (_io,)
+            R.output(gv1)
+        return gv1
+    # fmt: on
+
+    m = Model()
+    irmodule, _ = m.export_tvm(
+        spec={"test": {"x": spec.Tensor([1, 10], "float32")}},
+        debug=True,
+    )
+    tvm.ir.assert_structural_equal(irmodule["test"], test)
+
+
 def test_binary():
     class Model(Module):
         def test(self, x: Tensor, y: Tensor):
@@ -298,6 +323,7 @@ def test(self, x: Tensor, weight: Tensor, bias: Tensor):
             relu_out = op.relu(x)
             silu_out = op.silu(x)
             gelu_out = op.gelu(x)
+            sigmoid_out = op.sigmoid(x)
             softmax_out = op.softmax(x, axis=2)
             rms_norm_out = op.rms_norm(x, weight, axes=[-2, -1])
             rms_norm_with_bias_out = op.rms_norm(x, weight, axes=[-2, -1])
@@ -316,6 +342,7 @@ def test(
             relu: R.Tensor((2, 3, 4, 5), dtype="float32") = R.nn.relu(x)
             silu: R.Tensor((2, 3, 4, 5), dtype="float32") = R.nn.silu(x)
             gelu: R.Tensor((2, 3, 4, 5), dtype="float32") = R.nn.gelu(x)
+            sigmoid: R.Tensor((2, 3, 4, 5), dtype="float32") = R.sigmoid(x)
             softmax: R.Tensor((2, 3, 4, 5), dtype="float32") = R.nn.softmax(x, axis=2)
             rms_norm: R.Tensor((2, 3, 4, 5), dtype="float32") = R.nn.rms_norm(
                 x, weight, axes=[-2, -1], epsilon=1.0000000000000001e-05