tlopex commented on code in PR #20135:
URL: https://github.com/apache/tvm/pull/20135#discussion_r3858795108
##########
python/tvm/relax/frontend/torch/base_fx_graph_translator.py:
##########
@@ -572,23 +572,69 @@ def call_binary_op(op, lhs, rhs):
def _pow(self, node: fx.Node) -> relax.Var:
lhs, rhs = self.retrieve_args(node)
- # torch integer pow returns an integer tensor, but relax.op.power
legalizes to
- # TOPI power which requires floating-point inputs. Decompose an
integer base with
- # a constant non-negative integer exponent into repeated
multiplication instead.
- if (
- isinstance(lhs, relax.Expr)
- and isinstance(lhs.ty, relax.TensorType)
- and lhs.ty.dtype.matches_code(DataTypeCode.INT, DataTypeCode.UINT)
- and isinstance(rhs, int)
- and not isinstance(rhs, bool)
- and rhs >= 0
- ):
- if rhs == 0:
- return self.block_builder.emit(relax.op.ones_like(lhs))
- result = lhs
- for _ in range(rhs - 1):
- result = self.block_builder.emit(relax.op.multiply(result,
lhs))
- return result
+ if isinstance(lhs, relax.Expr) and isinstance(lhs.ty,
relax.TensorType):
+ lhs_dtype = lhs.ty.dtype
+ is_integer_base = lhs_dtype.matches_code(DataTypeCode.INT,
DataTypeCode.UINT)
+ is_float_base = lhs_dtype.matches_code(DataTypeCode.FLOAT,
DataTypeCode.BFLOAT)
+
+ # A Python float promotes an integer tensor to PyTorch's default
floating-point
+ # dtype. ExportedProgram records the inferred dtype, while plain
FX does not.
+ if is_integer_base and isinstance(rhs, float):
+ output_meta = node.meta.get("val")
+ output_dtype = self._convert_data_type(
+ output_meta.dtype
+ if isinstance(output_meta, self.torch.Tensor)
+ else self.torch.get_default_dtype()
+ )
+ lhs = self.block_builder.emit(relax.op.astype(lhs,
output_dtype))
+ lhs_dtype = lhs.ty.dtype
+ is_integer_base = False
+ is_float_base = True
+
+ # Match the scalar conversion used by PyTorch's floating-point
power kernels.
+ exponent_dtype = {
+ "float16": self.torch.float16,
+ "bfloat16": self.torch.bfloat16,
+ "float32": self.torch.float64,
+ "float64": self.torch.float64,
+ }.get(str(lhs_dtype))
+ if (
+ is_float_base
+ and exponent_dtype is not None
+ and isinstance(rhs, int | float)
+ and not isinstance(rhs, bool)
+ ):
+ rhs = self.torch.scalar_tensor(rhs, dtype=exponent_dtype,
device="cpu").item()
+
+ is_nonnegative_integral_exponent = (
+ isinstance(rhs, int) and not isinstance(rhs, bool) and rhs >= 0
+ ) or (isinstance(rhs, float) and rhs >= 0 and rhs.is_integer())
+
+ # TOPI power requires floating-point inputs, and some backends do
not preserve
+ # the sign of a negative base for integral exponents. Decompose
after applying
+ # PyTorch's dtype promotion so Python integer and float scalars
behave alike.
+ if (is_integer_base or is_float_base) and
is_nonnegative_integral_exponent:
+ exponent = int(rhs)
+ if exponent == 0:
+ return self.block_builder.emit(relax.op.ones_like(lhs))
+
+ # Exponentiation by squaring avoids linear graph growth for
large exponents.
Review Comment:
The dtype and exponent-conversion handling looks reasonable, but decomposing
floating-point `pow` into repeated multiplications changes PyTorch’s numerical
semantics because every intermediate `float16`/`bfloat16` multiplication is
rounded. e.g. `float16(0.361) ** 17` is nonzero in PyTorch but becomes zero in
TVM, while `bfloat16(1.1) ** 100` produces `15872` in PyTorch and `14848` in
TVM. Please preserve the `pow` kernel’s effective precision when fixing
negative-base behavior and add numerical tests with nontrivial low-precision
inputs; the current `-1` tests only verify exponent parity.
--
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]
