junrushao commented on code in PR #16316:
URL: https://github.com/apache/tvm/pull/16316#discussion_r1438835334
##########
python/tvm/relax/frontend/nn/core.py:
##########
@@ -274,6 +279,22 @@ def to(self, dtype: Optional[str] = None) -> None: #
pylint: disable=invalid-na
)._expr
+class Object:
+ """A wrapper on top of relax.Expr whose struct_info is the base
+ ObjectStructInfo (rather than any its subclass). Object effectively
+ represents non-tensor frontend components such as KV caches.
+ """
+
+ _expr: rx.Var
+
+ def __init__(self, *, _expr: rx.Expr, _name: str) -> None:
+ """Private constructor. Object is never supposed to be constructed
directly by users."""
+ if not isinstance(_expr, rx.Var):
+ _expr = BlockBuilder.current().emit(_expr, _name)
Review Comment:
got it! that's actually fine as it is, no need to change anything!
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
