diff --git a/python/tvm/relax/frontend/nn/op.py b/python/tvm/relax/frontend/nn/op.py index ac5858d5cd47..720e3dd3b429 100644 --- a/python/tvm/relax/frontend/nn/op.py +++ b/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, diff --git a/tests/python/relax/test_frontend_nn_op.py b/tests/python/relax/test_frontend_nn_op.py index 43f4a9efc03f..ed2e3753b2fb 100644 --- a/tests/python/relax/test_frontend_nn_op.py +++ b/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