gemini-code-assist[bot] commented on code in PR #18629:
URL: https://github.com/apache/tvm/pull/18629#discussion_r2655523646
##########
src/relax/op/nn/convolution.cc:
##########
@@ -707,14 +707,62 @@ StructInfo InferStructInfoConv1dTranspose(const Call&
call, const BlockBuilder&
return TensorStructInfo(ShapeExpr(out_shape), out_dtype, vdevice);
}
-// TODO(relax-team): implement FInferMixedPrecision and FRelaxInferLayout for
conv1d_transpose
-// and unit test for mixed_precision
+InferLayoutOutput InferLayoutConv1dTranspose(const Call& call,
+ const ffi::Map<ffi::String,
ffi::Array<ffi::String>>& desired_layouts,
+ const VarLayoutMap&
var_layout_map) {
+ const auto* attrs = call->attrs.as<Conv1DTransposeAttrs>();
+ LayoutDecision data_layout, weight_layout, output_layout;
+ ObjectPtr<Conv1DTransposeAttrs> new_attrs =
ffi::make_object<Conv1DTransposeAttrs>(*attrs);
+
+ auto it = desired_layouts.find("relax.nn.conv1d_transpose");
+ if (it != desired_layouts.end()) {
+ Layout desired_data_layout = (*it).second[0];
+ Layout desired_weight_layout = (*it).second[1];
+ Layout desired_output_layout = (*it).second.size() == 3 ? (*it).second[2]
: (*it).second[0];
Review Comment:
There is repeated code for accessing desired layouts. To improve readability
and maintainability, you can store the desired layout strings in local
variables and reuse them. This avoids repeating `(*it).second` access and the
ternary logic for the output layout. You can then use these variables on lines
730-732.
##########
src/relax/op/nn/convolution.cc:
##########
@@ -857,14 +905,95 @@ StructInfo InferStructInfoConv2dTranspose(const Call&
call, const BlockBuilder&
return TensorStructInfo(ShapeExpr(out_shape), out_dtype, vdevice);
}
-// TODO(relax-team): implement FInferMixedPrecision and FRelaxInferLayout for
conv2d_transpose
-// and unit test for mixed_precision
+InferLayoutOutput InferLayoutConv2dTranspose(const Call& call,
+ const ffi::Map<ffi::String,
ffi::Array<ffi::String>>& desired_layouts,
+ const VarLayoutMap&
var_layout_map) {
+ const auto* attrs = call->attrs.as<Conv2DTransposeAttrs>();
+ LayoutDecision data_layout = GetLayoutDecision(var_layout_map,
call->args[0]);
+ LayoutDecision weight_layout = GetLayoutDecision(var_layout_map,
call->args[1]);
+ LayoutDecision output_layout;
+ ObjectPtr<Conv2DTransposeAttrs> new_attrs =
ffi::make_object<Conv2DTransposeAttrs>(*attrs);
+
+ auto it = desired_layouts.find("relax.nn.conv2d_transpose");
+ if (it != desired_layouts.end()) {
+ Layout desired_data_layout = (*it).second[0];
+ Layout desired_weight_layout = (*it).second[1];
+ Layout desired_output_layout = (*it).second.size() == 3 ? (*it).second[2]
: (*it).second[0];
Review Comment:
There is repeated code for accessing desired layouts. To improve readability
and maintainability, you can store the desired layout strings in local
variables at the beginning of the `if` block and reuse them. This avoids
repeating `(*it).second` access and the ternary logic for the output layout.
You can then use these variables to update `new_attrs` in lines 933-935 and
957-959.
##########
tests/python/relax/test_op_nn_convolution.py:
##########
@@ -1571,6 +1590,25 @@ def
test_conv2d_transpose_infer_struct_info_wrong_input_type():
bb.normalize(relax.op.nn.conv2d_transpose(x1, w0))
+def test_conv2d_transpose_infer_struct_info_mixed_precision():
+ bb = relax.BlockBuilder()
+ x0 = relax.Var("x", R.Tensor((2, 3, 28, 28), "float16"))
+ w0 = relax.Var("w", R.Tensor((3, 4, 3, 3), "float16"))
+ x1 = relax.Var("x", R.Tensor((2, 3, 28, 28), "int8"))
+ w1 = relax.Var("w", R.Tensor((3, 4, 3, 3), "int8"))
+
+ _check_inference(
+ bb,
+ relax.op.nn.conv2d_transpose(x0, w0, out_dtype="float32"),
+ relax.TensorStructInfo((2, 4, 30, 30), "float32"),
+ )
+ _check_inference(
+ bb,
+ relax.op.nn.conv2d_transpose(x1, w1, out_dtype="int32"),
+ relax.TensorStructInfo((2, 4, 30, 30), "int32"),
+ )
Review Comment:
For better readability, consider using more descriptive variable names that
indicate the data type, like `x_f16`, `w_f16`, `x_i8`, and `w_i8`.
```suggestion
bb = relax.BlockBuilder()
x_f16 = relax.Var("x", R.Tensor((2, 3, 28, 28), "float16"))
w_f16 = relax.Var("w", R.Tensor((3, 4, 3, 3), "float16"))
x_i8 = relax.Var("x", R.Tensor((2, 3, 28, 28), "int8"))
w_i8 = relax.Var("w", R.Tensor((3, 4, 3, 3), "int8"))
_check_inference(
bb,
relax.op.nn.conv2d_transpose(x_f16, w_f16, out_dtype="float32"),
relax.TensorStructInfo((2, 4, 30, 30), "float32"),
)
_check_inference(
bb,
relax.op.nn.conv2d_transpose(x_i8, w_i8, out_dtype="int32"),
relax.TensorStructInfo((2, 4, 30, 30), "int32"),
)
```
##########
tests/python/relax/test_op_nn_convolution.py:
##########
@@ -782,6 +782,25 @@ def
test_conv1d_transpose_infer_struct_info_wrong_input_type():
bb.normalize(relax.op.nn.conv1d_transpose(x1, w0))
+def test_conv1d_transpose_infer_struct_info_mixed_precision():
+ bb = relax.BlockBuilder()
+ x0 = relax.Var("x", R.Tensor((2, 3, 28), "float16"))
+ w0 = relax.Var("w", R.Tensor((3, 4, 3), "float16"))
+ x1 = relax.Var("x", R.Tensor((2, 3, 28), "int8"))
+ w1 = relax.Var("w", R.Tensor((3, 4, 3), "int8"))
+
+ _check_inference(
+ bb,
+ relax.op.nn.conv1d_transpose(x0, w0, out_dtype="float32"),
+ relax.TensorStructInfo((2, 4, 30), "float32"),
+ )
+ _check_inference(
+ bb,
+ relax.op.nn.conv1d_transpose(x1, w1, out_dtype="int32"),
+ relax.TensorStructInfo((2, 4, 30), "int32"),
+ )
Review Comment:
For better readability, consider using more descriptive variable names that
indicate the data type, like `x_f16`, `w_f16`, `x_i8`, and `w_i8`.
```suggestion
bb = relax.BlockBuilder()
x_f16 = relax.Var("x", R.Tensor((2, 3, 28), "float16"))
w_f16 = relax.Var("w", R.Tensor((3, 4, 3), "float16"))
x_i8 = relax.Var("x", R.Tensor((2, 3, 28), "int8"))
w_i8 = relax.Var("w", R.Tensor((3, 4, 3), "int8"))
_check_inference(
bb,
relax.op.nn.conv1d_transpose(x_f16, w_f16, out_dtype="float32"),
relax.TensorStructInfo((2, 4, 30), "float32"),
)
_check_inference(
bb,
relax.op.nn.conv1d_transpose(x_i8, w_i8, out_dtype="int32"),
relax.TensorStructInfo((2, 4, 30), "int32"),
)
```
--
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]
