mirror of
https://github.com/ml-explore/mlx.git
synced 2025-10-20 17:38:09 +08:00
Added output_padding parameters in conv_transpose (#2092)
This commit is contained in:
@@ -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:
|
||||
|
Reference in New Issue
Block a user