gemini-code-assist[bot] commented on code in PR #19852:
URL: https://github.com/apache/tvm/pull/19852#discussion_r3447738609
##########
python/tvm/relax/script/parser/dist.py:
##########
@@ -34,33 +34,33 @@
from tvm.script.ir_builder.ir import IRModuleFrame
from tvm.tirx import PrimExpr
-from .entry import StructInfoProxy, TensorProxy
+from .entry import TensorProxy, TypeProxy
############################### R.DTensor ###############################
-class DTensorProxy(StructInfoProxy):
- tensor_sinfo_proxy: TensorProxy
+class DTensorProxy(TypeProxy):
+ tensor_ty_proxy: TensorProxy
device_mesh: DeviceMesh
placement: Placement
def __init__(
self,
- tensor_sinfo_proxy: TensorProxy,
+ tensor_ty_proxy: TensorProxy,
device_mesh: DeviceMesh,
placement: Placement,
) -> None:
self.device_mesh = device_mesh
self.placement = placement
- self.tensor_sinfo_proxy = tensor_sinfo_proxy
+ self.tensor_ty_proxy = tensor_ty_proxy
super().__init__()
def get_symbolic_vars(self) -> set[str]:
- return self.tensor_sinfo_proxy.get_symbolic_vars()
+ return self.tensor_ty_proxy.get_symbolic_vars()
- def as_struct_info(self, dict_globals: dict[str, Any] | None = None) ->
TensorStructInfo:
- return DTensorStructInfo(
- self.tensor_sinfo_proxy.as_struct_info(dict_globals),
+ def as_ty(self, dict_globals: dict[str, Any] | None = None) -> TensorType:
Review Comment:

The return type annotation of `as_ty` in `DTensorProxy` is incorrectly
specified as `TensorType` instead of `DTensorType`. Since `DTensorType`
inherits from `Type` (not `TensorType`), this is a type annotation mismatch.
```suggestion
def as_ty(self, dict_globals: dict[str, Any] | None = None) ->
DTensorType:
```
##########
src/relax/backend/adreno/annotate_custom_storage.cc:
##########
@@ -357,12 +357,11 @@ class CollectConsumerScopeInfo : public ExprVisitor {
ffi::Array<ffi::String> arg_scope;
for (uint32_t i = 0; i < func_args->fields.size(); ++i) {
- auto sinfo = GetStructInfo(func_args->fields[i]);
- if (auto tensor_sinfo = sinfo.as<TensorStructInfo>()) {
+ auto ty = GetType(func_args->fields[i]);
+ if (auto tensor_ty = ty.as<TensorType>()) {
bool is_texture =
i < is_texture_supported.size() ? is_texture_supported[i] :
is_texture_supported[0];
- auto scope =
- is_texture ?
Scope(GetShapeFromTensorStructInfo(tensor_sinfo.value())) : "global";
+ auto scope = is_texture ?
Scope(GetShapeFromTensorType(tensor_ty.value())) : "global";
Review Comment:

Using `ty.as<TensorType>()` will result in a compilation error because
`TensorType` is an `ObjectRef` subclass, not an `Object` subclass. It should be
`ty.as<TensorTypeNode>()`. Additionally, `tensor_ty` is a pointer (`const
TensorTypeNode*`), so calling `.value()` on it is invalid. Instead, use
`ffi::GetRef<TensorType>(tensor_ty)` to pass the reference to
`GetShapeFromTensorType`.
```suggestion
if (auto tensor_ty = ty.as<TensorTypeNode>()) {
bool is_texture =
i < is_texture_supported.size() ? is_texture_supported[i] :
is_texture_supported[0];
auto scope = is_texture ?
Scope(GetShapeFromTensorType(ffi::GetRef<TensorType>(tensor_ty))) : "global";
```
##########
python/tvm/relax/block_builder.py:
##########
@@ -114,8 +114,8 @@ class BlockBuilder(Object):
m = tirx.Var("m", "int32")
n = tirx.Var("n", "int32")
- x = rx.Var("x", rx.TensorStructInfo([m, n], "float16"))
- y = rx.Var("y", rx.TensorStructInfo([n], "float16")
+ x = rx.Var("x", rx.TensorType([m, n], "float16"))
+ y = rx.Var("y", rx.TensorType([n], "float16")
Review Comment:

There is a missing closing parenthesis `)` at the end of the line in the
docstring example.
```suggestion
y = rx.Var("y", rx.TensorType([n], "float16"))
```
##########
python/tvm/relax/script/parser/entry.py:
##########
@@ -478,12 +476,12 @@ def get_symbolic_vars(self) -> set[str]:
else:
return set()
- def as_struct_info(self, dict_globals: dict[str, Any] | None = None) ->
ShapeStructInfo:
+ def as_ty(self, dict_globals: dict[str, Any] | None = None) -> ShapeType:
Review Comment:

The return type annotation of `as_ty` in `PrimProxy` is incorrectly
specified as `ShapeType` instead of `PrimType`. This appears to be a copy-paste
error from `ShapeProxy`.
```suggestion
def as_ty(self, dict_globals: dict[str, Any] | None = None) -> PrimType:
```
##########
python/tvm/relax/script/parser/entry.py:
##########
@@ -177,8 +177,8 @@ def __init__(self) -> None:
def get_symbolic_vars(self) -> set[str]:
return set()
- def as_struct_info(self, dict_globals: dict[str, Any] | None = None) ->
ShapeStructInfo:
- return ObjectStructInfo()
+ def as_ty(self, dict_globals: dict[str, Any] | None = None) -> ShapeType:
Review Comment:

The return type annotation of `as_ty` in `ObjectProxy` is incorrectly
specified as `ShapeType` instead of `ObjectType`. This appears to be a
copy-paste error from `ShapeProxy`.
```suggestion
def as_ty(self, dict_globals: dict[str, Any] | None = None) ->
ObjectType:
```
--
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]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]