[Unity][Frontend] Add Sqrt Op

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

@@ -17,7 +17,7 @@
 """The core infra for nn.Module, which includes the following pieces:
 - Tensor, a wrapper on top of relax.Expr whose struct_info is a TensorStructInfo,
   providing more convenient access shape and dtype information.
-  Tensor is always symbolc and not bound to any concrete values.
+  Tensor is always symbolic and not bound to any concrete values.
 - Parameter, a special tensor which could be bound or not bound to concrete values.
 - Module, a container of nn.Parameters and sub nn.Modules.
 - Effect, a non-user-facing class that encloses potential side effects, for example, IO,

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

@@ -1486,6 +1486,28 @@ def square(x: Tensor, name: str = "square") -> Tensor:
     return wrap_nested(_op.square(x._expr), name)
 
 
+def sqrt(x: Tensor, name: str = "sqrt") -> Tensor:
+    """Computes the element-wise sqrt of the input tensor.
+
+    Parameters
+    ----------
+    x : Tensor
+        The input tensor.
+
+    name : str
+        Name hint.
+
+    Returns
+    -------
+    result : Tensor
+        The computed result.
+    Note
+    ----
+    The input tensor is required to have float dtype
+    """
+    return wrap_nested(_op.sqrt(x._expr), name)
+
+
 def get_timestep_embedding(
     x: Tensor,
     embedding_dim: int,

tests/python/relax/test_frontend_nn_op.py

@@ -31,15 +31,17 @@ def test_unary():
     class Model(Module):
         def test(self, x: Tensor):
             z0 = op.square(x)
-            return (x,)
+            z1 = op.sqrt(x)
+            return (z0, z1)
 
     # 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,)
+            sqrt: R.Tensor((1, 10), dtype="float32") = R.sqrt(x)
+            gv1 = (square, sqrt), (_io,)
             R.output(gv1)
         return gv1
     # fmt: on