Added output_padding parameters in conv_transpose (#2092)

This commit is contained in:
Param Thakkar
2025-04-23 21:56:33 +05:30
committed by GitHub
parent 3836445241
commit 600e87e03c
6 changed files with 366 additions and 14 deletions

View File

@@ -3609,11 +3609,12 @@ void init_ops(nb::module_& m) {
"stride"_a = 1,
"padding"_a = 0,
"dilation"_a = 1,
"output_padding"_a = 0,
"groups"_a = 1,
nb::kw_only(),
"stream"_a = nb::none(),
nb::sig(
"def conv_transpose1d(input: array, weight: array, /, stride: int = 1, padding: int = 0, dilation: int = 1, groups: int = 1, *, stream: Union[None, Stream, Device] = None) -> array"),
"def conv_transpose1d(input: array, weight: array, /, stride: int = 1, padding: int = 0, dilation: int = 1, output_padding: int = 0, groups: int = 1, *, stream: Union[None, Stream, Device] = None) -> array"),
R"pbdoc(
1D transposed convolution over an input with several channels
@@ -3623,6 +3624,7 @@ void init_ops(nb::module_& m) {
stride (int, optional): Kernel stride. Default: ``1``.
padding (int, optional): Input padding. Default: ``0``.
dilation (int, optional): Kernel dilation. Default: ``1``.
output_padding (int, optional): Output padding. Default: ``0``.
groups (int, optional): Input feature groups. Default: ``1``.
Returns:
@@ -3635,11 +3637,13 @@ void init_ops(nb::module_& m) {
const std::variant<int, std::pair<int, int>>& stride,
const std::variant<int, std::pair<int, int>>& padding,
const std::variant<int, std::pair<int, int>>& dilation,
const std::variant<int, std::pair<int, int>>& output_padding,
int groups,
mx::StreamOrDevice s) {
std::pair<int, int> stride_pair{1, 1};
std::pair<int, int> padding_pair{0, 0};
std::pair<int, int> dilation_pair{1, 1};
std::pair<int, int> output_padding_pair{0, 0};
if (auto pv = std::get_if<int>(&stride); pv) {
stride_pair = std::pair<int, int>{*pv, *pv};
@@ -3659,19 +3663,33 @@ void init_ops(nb::module_& m) {
dilation_pair = std::get<std::pair<int, int>>(dilation);
}
if (auto pv = std::get_if<int>(&output_padding); pv) {
output_padding_pair = std::pair<int, int>{*pv, *pv};
} else {
output_padding_pair = std::get<std::pair<int, int>>(output_padding);
}
return mx::conv_transpose2d(
input, weight, stride_pair, padding_pair, dilation_pair, groups, s);
input,
weight,
stride_pair,
padding_pair,
dilation_pair,
output_padding_pair,
groups,
s);
},
nb::arg(),
nb::arg(),
"stride"_a = 1,
"padding"_a = 0,
"dilation"_a = 1,
"output_padding"_a = 0,
"groups"_a = 1,
nb::kw_only(),
"stream"_a = nb::none(),
nb::sig(
"def conv_transpose2d(input: array, weight: array, /, stride: Union[int, Tuple[int, int]] = 1, padding: Union[int, Tuple[int, int]] = 0, dilation: Union[int, Tuple[int, int]] = 1, groups: int = 1, *, stream: Union[None, Stream, Device] = None) -> array"),
"def conv_transpose2d(input: array, weight: array, /, stride: Union[int, Tuple[int, int]] = 1, padding: Union[int, Tuple[int, int]] = 0, dilation: Union[int, Tuple[int, int]] = 1, output_padding: Union[int, Tuple[int, int]] = 0, groups: int = 1, *, stream: Union[None, Stream, Device] = None) -> array"),
R"pbdoc(
2D transposed convolution over an input with several channels
@@ -3689,6 +3707,9 @@ void init_ops(nb::module_& m) {
dilation (int or tuple(int), optional): :obj:`tuple` of size 2 with
kernel dilation. All spatial dimensions get the same dilation
if only one number is specified. Default: ``1``
output_padding (int or tuple(int), optional): :obj:`tuple` of size 2 with
output padding. All spatial dimensions get the same output
padding if only one number is specified. Default: ``0``.
groups (int, optional): input feature groups. Default: ``1``.
Returns:
@@ -3701,11 +3722,13 @@ void init_ops(nb::module_& m) {
const std::variant<int, std::tuple<int, int, int>>& stride,
const std::variant<int, std::tuple<int, int, int>>& padding,
const std::variant<int, std::tuple<int, int, int>>& dilation,
const std::variant<int, std::tuple<int, int, int>>& output_padding,
int groups,
mx::StreamOrDevice s) {
std::tuple<int, int, int> stride_tuple{1, 1, 1};
std::tuple<int, int, int> padding_tuple{0, 0, 0};
std::tuple<int, int, int> dilation_tuple{1, 1, 1};
std::tuple<int, int, int> output_padding_tuple{0, 0, 0};
if (auto pv = std::get_if<int>(&stride); pv) {
stride_tuple = std::tuple<int, int, int>{*pv, *pv, *pv};
@@ -3725,12 +3748,20 @@ void init_ops(nb::module_& m) {
dilation_tuple = std::get<std::tuple<int, int, int>>(dilation);
}
if (auto pv = std::get_if<int>(&output_padding); pv) {
output_padding_tuple = std::tuple<int, int, int>{*pv, *pv, *pv};
} else {
output_padding_tuple =
std::get<std::tuple<int, int, int>>(output_padding);
}
return mx::conv_transpose3d(
input,
weight,
stride_tuple,
padding_tuple,
dilation_tuple,
output_padding_tuple,
groups,
s);
},
@@ -3739,11 +3770,12 @@ void init_ops(nb::module_& m) {
"stride"_a = 1,
"padding"_a = 0,
"dilation"_a = 1,
"output_padding"_a = 0,
"groups"_a = 1,
nb::kw_only(),
"stream"_a = nb::none(),
nb::sig(
"def conv_transpose3d(input: array, weight: array, /, stride: Union[int, Tuple[int, int, int]] = 1, padding: Union[int, Tuple[int, int, int]] = 0, dilation: Union[int, Tuple[int, int, int]] = 1, groups: int = 1, *, stream: Union[None, Stream, Device] = None) -> array"),
"def conv_transpose3d(input: array, weight: array, /, stride: Union[int, Tuple[int, int, int]] = 1, padding: Union[int, Tuple[int, int, int]] = 0, dilation: Union[int, Tuple[int, int, int]] = 1, output_padding: Union[int, Tuple[int, int, int]] = 0, groups: int = 1, *, stream: Union[None, Stream, Device] = None) -> array"),
R"pbdoc(
3D transposed convolution over an input with several channels
@@ -3761,6 +3793,9 @@ void init_ops(nb::module_& m) {
dilation (int or tuple(int), optional): :obj:`tuple` of size 3 with
kernel dilation. All spatial dimensions get the same dilation
if only one number is specified. Default: ``1``
output_padding (int or tuple(int), optional): :obj:`tuple` of size 3 with
output padding. All spatial dimensions get the same output
padding if only one number is specified. Default: ``0``.
groups (int, optional): input feature groups. Default: ``1``.
Returns: