gemini-code-assist[bot] commented on code in PR #19499:
URL: https://github.com/apache/tvm/pull/19499#discussion_r3178244271
##########
python/tvm/relax/frontend/tflite/tflite_frontend.py:
##########
@@ -3282,8 +3282,26 @@ def convert_batch_to_space_nd(self, op):
block_shape = list(self.get_tensor_value(input_tensors[1]))
crops = self.get_tensor_value(input_tensors[2]).tolist()
+ crop_begin = [crop[0] for crop in crops]
+ crop_end = [crop[1] for crop in crops]
- out = relax.op.nn.batch_to_space_nd(in_expr, block_shape, crops)
+ output_tensors = self.get_output_tensors(op)
+ assert len(output_tensors) == 1, "output tensors length should be 1"
+ output_tensor = output_tensors[0]
+ output_shape = to_int_list(self.get_tensor_shape(output_tensor))
+ output_dtype = self.get_tensor_type_str(output_tensor.tensor.Type())
+
+ out = relax.op.call_dps_packed(
+ "topi.nn.batch_to_space_nd",
+ (
+ in_expr,
+ relax.ShapeExpr(block_shape),
+ relax.ShapeExpr(crop_begin),
+ relax.ShapeExpr(crop_end),
+ "batch_to_space_nd",
+ ),
+ out_sinfo=relax.TensorStructInfo(output_shape, output_dtype),
+ )
Review Comment:
There are two issues in `convert_batch_to_space_nd`:
1. **Argument Count Bug**: The `call_dps_packed` call for
`topi.nn.batch_to_space_nd` includes a 5th argument `"batch_to_space_nd"`.
However, the TOPI implementation of `batch_to_space_nd` only accepts 4
arguments: `data`, `block_shape`, `crop_begin_list`, and `crop_end_list`. This
will cause a runtime error during execution.
2. **Type Safety**: It is recommended to use `to_int_list` for
`block_shape`, `crop_begin`, and `crop_end`. This ensures that numpy integer
types are converted to Python integers, preventing potential type mismatch
issues (e.g., being incorrectly lowered as `int64` constants) in the Relax IR.
```python
block_shape = to_int_list(self.get_tensor_value(input_tensors[1]))
crops = self.get_tensor_value(input_tensors[2])
crop_begin = to_int_list(crops[:, 0])
crop_end = to_int_list(crops[:, 1])
output_tensors = self.get_output_tensors(op)
assert len(output_tensors) == 1, "output tensors length should be 1"
output_tensor = output_tensors[0]
output_shape = to_int_list(self.get_tensor_shape(output_tensor))
output_dtype = self.get_tensor_type_str(output_tensor.tensor.Type())
out = relax.op.call_dps_packed(
"topi.nn.batch_to_space_nd",
(
in_expr,
relax.ShapeExpr(block_shape),
relax.ShapeExpr(crop_begin),
relax.ShapeExpr(crop_end),
),
out_sinfo=relax.TensorStructInfo(output_shape, output_dtype),
)
```
##########
python/tvm/relax/frontend/tflite/tflite_frontend.py:
##########
@@ -3391,8 +3409,26 @@ def convert_space_to_batch_nd(self, op):
block_shape = list(self.get_tensor_value(input_tensors[1]))
paddings = self.get_tensor_value(input_tensors[2]).tolist()
+ pad_before = [pad[0] for pad in paddings]
+ pad_after = [pad[1] for pad in paddings]
- out = relax.op.nn.space_to_batch_nd(in_expr, block_shape, paddings)
+ output_tensors = self.get_output_tensors(op)
+ assert len(output_tensors) == 1, "output tensors length should be 1"
+ output_tensor = output_tensors[0]
+ output_shape = to_int_list(self.get_tensor_shape(output_tensor))
+ output_dtype = self.get_tensor_type_str(output_tensor.tensor.Type())
+
+ out = relax.op.call_dps_packed(
+ "topi.nn.space_to_batch_nd",
+ (
+ in_expr,
+ relax.ShapeExpr(block_shape),
+ relax.ShapeExpr(pad_before),
+ relax.ShapeExpr(pad_after),
+ 0.0,
+ ),
+ out_sinfo=relax.TensorStructInfo(output_shape, output_dtype),
+ )
Review Comment:
Similar to `convert_batch_to_space_nd`, please use `to_int_list` for
`block_shape`, `pad_before`, and `pad_after` to ensure proper type conversion
for Relax and avoid potential `IntImm` width issues.
```python
block_shape = to_int_list(self.get_tensor_value(input_tensors[1]))
paddings = self.get_tensor_value(input_tensors[2])
pad_before = to_int_list(paddings[:, 0])
pad_after = to_int_list(paddings[:, 1])
output_tensors = self.get_output_tensors(op)
assert len(output_tensors) == 1, "output tensors length should be 1"
output_tensor = output_tensors[0]
output_shape = to_int_list(self.get_tensor_shape(output_tensor))
output_dtype = self.get_tensor_type_str(output_tensor.tensor.Type())
out = relax.op.call_dps_packed(
"topi.nn.space_to_batch_nd",
(
in_expr,
relax.ShapeExpr(block_shape),
relax.ShapeExpr(pad_before),
relax.ShapeExpr(pad_after),
0.0,
),
out_sinfo=relax.TensorStructInfo(output_shape, output_dtype),
)
```
--
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]
