gemini-code-assist[bot] commented on code in PR #18670:
URL: https://github.com/apache/tvm/pull/18670#discussion_r2703035468
##########
python/tvm/relax/frontend/torch/exported_program_translator.py:
##########
@@ -1812,8 +1818,31 @@ def forward(self, input):
# Use the importer to import the ExportedProgram to Relax.
mod: tvm.IRModule = from_exported_program(exported_program)
"""
+ def _is_sparse_tensor(value: object) -> bool:
+ if not isinstance(value, torch.Tensor):
+ return False
+ try:
+ return value.layout != torch.strided
+ except RuntimeError:
+ return False
+
+ def _has_sparse_tensors(ep: torch.export.ExportedProgram) -> bool:
+ for _, tensor in ep.named_buffers():
+ if _is_sparse_tensor(tensor):
+ return True
+ for _, tensor in ep.named_parameters():
+ if _is_sparse_tensor(tensor):
+ return True
+ for tensor in ep.constants.values():
+ if _is_sparse_tensor(tensor):
+ return True
+ for tensor in ep.tensor_constants.values():
+ if _is_sparse_tensor(tensor):
+ return True
+ return False
Review Comment:
This function can be made more concise and readable by using
`itertools.chain` to combine the different tensor iterators and `any()` to
perform the check. This avoids repeating the loop structure and improves
maintainability.
```suggestion
def _has_sparse_tensors(ep: torch.export.ExportedProgram) -> bool:
from itertools import chain
all_potential_tensors = chain(
(t for _, t in ep.named_buffers()),
(t for _, t in ep.named_parameters()),
ep.constants.values(),
ep.tensor_constants.values(),
)
return any(_is_sparse_tensor(t) for t in all_potential_tensors)
```
--
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]
