MasterJH5574 commented on code in PR #16316:
URL: https://github.com/apache/tvm/pull/16316#discussion_r1438768456
##########
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)
+ self._expr = _expr
+ assert type(self._expr.struct_info) is ObjectStructInfo
Review Comment:
This is on purpose if we want to make sure the frontend `Object` can only
represent “opaque objects” and never express other types (like a tensor).
https://github.com/apache/tvm/blob/f671711e6c7e603a0ecc3621f7ea831beb3ade75/python/tvm/relax/frontend/nn/core.py#L282-L285
Say if we pass in a Var with TensorStructInfo here, `isinstance(sinfo,
ObjectStructInfo)` returns true, and `type(self._expr.struct_info) is
ObjectStructInfo` returns false.
--
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]
