javierdejesusda commented on code in PR #19660:
URL: https://github.com/apache/tvm/pull/19660#discussion_r3348513537
##########
python/tvm/relax/frontend/torch/base_fx_graph_translator.py:
##########
@@ -523,6 +524,27 @@ def call_binary_op(op, lhs, rhs):
return convert
+ 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.struct_info, relax.TensorStructInfo)
+ and "int" in lhs.struct_info.dtype
+ 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
Review Comment:
Keeping linear repeated multiplication intentionally: it matches the
existing frontend idiom (onnx_frontend.py:1627, x^3 = x*x*x), and because each
factor is emitted as a separate flat binding there's no recursive nesting to
overflow. Realistic integer exponents are small, and large ones overflow the
integer dtype itself, so O(log N) buys nothing here while complicating the
expected-IR tests.
--
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]
