gemini-code-assist[bot] commented on code in PR #19732:
URL: https://github.com/apache/tvm/pull/19732#discussion_r3394043393
##########
python/tvm/relax/frontend/torch/base_fx_graph_translator.py:
##########
@@ -1996,6 +2010,17 @@ def _index_tensor(self, node: fx.Node) -> relax.Var:
axis, index_tensor = non_none_indices[0]
return self.block_builder.emit(relax.op.take(data, index_tensor,
axis=axis))
+ # If no dimension is sliced (no None entries), this is plain
NumPy-style
+ # advanced indexing: the index tensors broadcast together and are
applied
+ # jointly ("zipped"), which is exactly relax.op.index_tensor's
semantics.
+ # Note that the sequential-take path below is NOT equivalent: it
computes
+ # an outer product over the index tensors, which only matches PyTorch
+ # when the index shapes are mutually orthogonal (e.g. (H, 1) and (W,)).
+ if len(non_none_indices) == len(indices):
+ return self.block_builder.emit(
+ relax.op.index_tensor(data, [idx for _, idx in
non_none_indices])
+ )
Review Comment:
If `indices` is empty, `len(non_none_indices) == len(indices)` will evaluate
to `True` (0 == 0). This will cause the translator to emit
`relax.op.index_tensor(data, [])` instead of returning `data` directly (which
was the behavior of the fallback path). To prevent emitting an empty
`index_tensor` op, we should ensure `len(indices) > 0` before applying this
optimization.
```python
if len(indices) > 0 and len(non_none_indices) == len(indices):\n
return self.block_builder.emit(\n
relax.op.index_tensor(data, [idx for _, idx in non_none_indices])\n )
```
##########
python/tvm/relax/frontend/torch/base_fx_graph_translator.py:
##########
@@ -1684,6 +1684,20 @@ def _sum(self, node: fx.Node) -> relax.Var:
if isinstance(dim, list | tuple) and len(dim) == 0:
dim = None
keepdim = args[2] if len(node.args) > 2 else
node.kwargs.get("keepdim", False)
+ dtype = node.kwargs.get("dtype", None)
+ if dtype is not None:
+ x = self.block_builder.emit(
+ relax.op.astype(x, self._convert_data_type(dtype, self.env))
+ )
+ else:
+ # Match PyTorch type promotion: summing bool or integer tensors
+ # accumulates in int64 unless an explicit dtype is given.
+ input_dtype = x.struct_info.dtype
+ if input_dtype == "bool" or (
+ (input_dtype.startswith("int") or
input_dtype.startswith("uint"))
+ and input_dtype != "int64"
+ ):
+ x = self.block_builder.emit(relax.op.astype(x, "int64"))
Review Comment:
There are two improvement opportunities here to fully align with PyTorch's
type promotion rules and ensure robust execution:\n\n1. **Defensive
Programming**: Accessing `x.struct_info.dtype` directly assumes `x.struct_info`
is always populated and is an instance of `relax.TensorStructInfo`. If it is
`None` or of another type (e.g., `TupleStructInfo`), this will raise an
`AttributeError`. We should guard this with `isinstance(x.struct_info,
relax.TensorStructInfo)`.\n2. **Float16/BFloat16 Promotion**: According to
PyTorch's `torch.sum` specification, `float16` and `bfloat16` inputs are
accumulated in `float32` by default to prevent overflow and precision loss. We
should promote these to `float32` when no explicit `dtype` is provided.
```python
if isinstance(x.struct_info, relax.TensorStructInfo):\n
input_dtype = x.struct_info.dtype\n if input_dtype ==
"bool" or (\n (input_dtype.startswith("int") or
input_dtype.startswith("uint"))\n and input_dtype !=
"int64"\n ):\n x =
self.block_builder.emit(relax.op.astype(x, "int64"))\n elif
input_dtype in ("float16", "bfloat16"):\n x =
self.block_builder.emit(relax.op.astype(x, "float32"))
```
--
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]
