gemini-code-assist[bot] commented on code in PR #19756:
URL: https://github.com/apache/tvm/pull/19756#discussion_r3408079906
##########
python/tvm/relax/frontend/torch/base_fx_graph_translator.py:
##########
@@ -410,6 +410,28 @@ def _logical_not(self, node: fx.Node) -> relax.Var:
x = self.block_builder.emit(relax.op.astype(x, "bool"))
return self.block_builder.emit(relax.op.logical_not(x))
+ def _logical_or(self, node: fx.Node) -> relax.Var:
+ lhs = self.env[node.args[0]]
+ rhs = self.env[node.args[1]]
+ # torch.logical_or accepts any dtype (treating nonzero as True) and
returns bool, but
+ # relax.op.logical_or requires boolean inputs, so cast non-bool inputs
to bool first.
+ if lhs.struct_info.dtype != "bool":
+ lhs = self.block_builder.emit(relax.op.astype(lhs, "bool"))
+ if rhs.struct_info.dtype != "bool":
+ rhs = self.block_builder.emit(relax.op.astype(rhs, "bool"))
+ return self.block_builder.emit(relax.op.logical_or(lhs, rhs))
+
+ def _logical_xor(self, node: fx.Node) -> relax.Var:
+ lhs = self.env[node.args[0]]
+ rhs = self.env[node.args[1]]
+ # torch.logical_xor accepts any dtype (treating nonzero as True) and
returns bool, but
+ # relax.op.logical_xor requires boolean inputs, so cast non-bool
inputs to bool first.
+ if lhs.struct_info.dtype != "bool":
+ lhs = self.block_builder.emit(relax.op.astype(lhs, "bool"))
+ if rhs.struct_info.dtype != "bool":
+ rhs = self.block_builder.emit(relax.op.astype(rhs, "bool"))
+ return self.block_builder.emit(relax.op.logical_xor(lhs, rhs))
Review Comment:
To improve maintainability and robustness, we can refactor the logical
binary operations into a common helper method. This avoids duplicating the
casting logic and allows us to safely access the `dtype` attribute using
`getattr` to prevent potential `AttributeError`s if `struct_info` is `None` or
is not a tensor struct info (e.g., `TupleStructInfo`).
```suggestion
def _logical_binary_op(self, node: fx.Node, op) -> relax.Var:
lhs = self.env[node.args[0]]
rhs = self.env[node.args[1]]
if getattr(lhs.struct_info, "dtype", None) != "bool":
lhs = self.block_builder.emit(relax.op.astype(lhs, "bool"))
if getattr(rhs.struct_info, "dtype", None) != "bool":
rhs = self.block_builder.emit(relax.op.astype(rhs, "bool"))
return self.block_builder.emit(op(lhs, rhs))
def _logical_or(self, node: fx.Node) -> relax.Var:
return self._logical_binary_op(node, relax.op.logical_or)
def _logical_xor(self, node: fx.Node) -> relax.Var:
return self._logical_binary_op(node, relax.op.logical_xor)
```
--
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]
