(tvm) branch unity updated: [Unity][Frontend] Add the `sum` op to frontend ops (#16225)

[tq...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

tqchen pushed a commit to branch unity
in repository https://gitbox.apache.org/repos/asf/tvm.git


The following commit(s) were added to refs/heads/unity by this push:
     new cbcb67c047 [Unity][Frontend] Add the `sum` op to frontend ops (#16225)
cbcb67c047 is described below

commit cbcb67c047126ca725cc529fcb7533a4a3e0d694
Author: Ruihang Lai <ru...@cs.cmu.edu>
AuthorDate: Wed Dec 13 16:37:48 2023 -0500

    [Unity][Frontend] Add the `sum` op to frontend ops (#16225)
---
 python/tvm/relax/block_builder.py         |  2 +-
 python/tvm/relax/frontend/nn/__init__.py  |  1 +
 python/tvm/relax/frontend/nn/op.py        | 35 +++++++++++++++++++++++++++++++
 tests/python/relax/test_frontend_nn_op.py | 32 +++++++++++++++++++++-------
 4 files changed, 62 insertions(+), 8 deletions(-)

diff --git a/python/tvm/relax/block_builder.py b/python/tvm/relax/block_builder.py
index ab807203a3..0a408578df 100644
--- a/python/tvm/relax/block_builder.py
+++ b/python/tvm/relax/block_builder.py
@@ -152,7 +152,7 @@ class BlockBuilder(Object):
     _stack = []
 
     @staticmethod
-    def current():
+    def current() -> Optional["BlockBuilder"]:
         """Returns the current BlockBuilder."""
         if BlockBuilder._stack:
             return BlockBuilder._stack[-1]
diff --git a/python/tvm/relax/frontend/nn/__init__.py b/python/tvm/relax/frontend/nn/__init__.py
index 4e11f167b5..a2c6f39442 100644
--- a/python/tvm/relax/frontend/nn/__init__.py
+++ b/python/tvm/relax/frontend/nn/__init__.py
@@ -15,6 +15,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """A PyTorch-like API to build IRModules."""
+# pylint: disable=redefined-builtin
 from . import op, spec
 from .core import (
     Effect,
diff --git a/python/tvm/relax/frontend/nn/op.py b/python/tvm/relax/frontend/nn/op.py
index 76af0044c3..d315331862 100644
--- a/python/tvm/relax/frontend/nn/op.py
+++ b/python/tvm/relax/frontend/nn/op.py
@@ -15,6 +15,7 @@
 # specific language governing permissions and limitations
 # under the License.
 # pylint: disable=too-many-lines,invalid-name,protected-access,redefined-outer-name
+# pylint: disable=redefined-builtin
 """nn.Tensor operators."""
 import math
 from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, Union
@@ -241,6 +242,40 @@ def chunk(x: Tensor, chunks: int, dim: int = 0, name: str = "chunk") -> Tensor:
     return _wrap_nested(_op.split(x._expr, chunks, dim), name)
 
 
+def sum(
+    x: Tensor,
+    axis: Optional[Union[int, List[int]]] = None,
+    keepdims: bool = False,
+    name: str = "sum",
+) -> Tensor:
+    """Computes the sum of tensor elements over given axes.
+
+    Parameters
+    ----------
+    x : Tensor
+        The input data tensor
+
+    axis : Optional[Union[int, List[int]]]
+        Axis or axes along which a sum is performed.
+        The default, axis=None, will sum all of the elements of the input tensor.
+        Negative indexing is supported.
+
+    keepdims : bool
+        If this is set to True, the axes which are reduced are left in the result as
+        dimensions with size one.
+        With this option, the result will broadcast correctly against the input tensor.
+
+    name : str
+        Name hint for this operation.
+
+    Returns
+    -------
+    result : Tensor
+        The computed result.
+    """
+    return _wrap_nested(_op.sum(x._expr, axis, keepdims), name)
+
+
 def matmul(a: Tensor, b: Tensor, out_dtype: Optional[str] = None, name: str = "matmul") -> Tensor:
     """General matrix multiplication of two tensors, with broadcasting on batched dimensions.
 
diff --git a/tests/python/relax/test_frontend_nn_op.py b/tests/python/relax/test_frontend_nn_op.py
index 57bf6d273c..e952316376 100644
--- a/tests/python/relax/test_frontend_nn_op.py
+++ b/tests/python/relax/test_frontend_nn_op.py
@@ -14,15 +14,9 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
-import io
-import sys
-
-import pytest
-import torch
-
 import tvm
 import tvm.testing
-from tvm import relax, te, tir
+from tvm import tir
 from tvm.relax.frontend.nn import Module, Tensor, op, spec
 from tvm.script import ir as I
 from tvm.script import relax as R
@@ -66,6 +60,30 @@ def test_binary():
     tvm.ir.assert_structural_equal(irmodule["test"], test)
 
 
+def test_sum():
+    class Model(Module):
+        def test(self, x: Tensor):
+            z0 = op.sum(x, axis=[1, 2], keepdims=True)
+            return z0
+
+    # fmt: off
+    @R.function
+    def test(x: R.Tensor((3, 5, 2, 4), dtype="float32"), _io: R.Object) -> R.Tuple(R.Tensor((3, 1, 1, 4), dtype="float32"), R.Tuple(R.Object)):
+        R.func_attr({"num_input": 2})
+        with R.dataflow():
+            sum: R.Tensor((3, 1, 1, 4), dtype="float32") = R.sum(x, axis=[1, 2], keepdims=True)
+            gv1: R.Tuple(R.Tensor((3, 1, 1, 4), dtype="float32"), R.Tuple(R.Object)) = sum, (_io,)
+            R.output(gv1)
+        return gv1
+    # fmt: on
+
+    m = Model()
+    irmodule, _ = m.export_tvm(
+        spec={"test": {"x": spec.Tensor([3, 5, 2, 4], "float32")}}, debug=True
+    )
+    tvm.ir.assert_structural_equal(irmodule["test"], test)
+
+
 def test_manipulate():
     class Model(Module):
         def test(self, x: Tensor):