xref: /aosp_15_r20/external/pytorch/torch/csrc/api/include/torch/nn/functional/pixelshuffle.h (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
1 #pragma once
2 
3 #include <torch/nn/options/pixelshuffle.h>
4 
5 namespace torch {
6 namespace nn {
7 namespace functional {
8 
9 #ifndef DOXYGEN_SHOULD_SKIP_THIS
10 namespace detail {
pixel_shuffle(const Tensor & input,int64_t upscale_factor)11 inline Tensor pixel_shuffle(const Tensor& input, int64_t upscale_factor) {
12   return torch::pixel_shuffle(input, upscale_factor);
13 }
14 
pixel_unshuffle(const Tensor & input,int64_t downscale_factor)15 inline Tensor pixel_unshuffle(const Tensor& input, int64_t downscale_factor) {
16   return torch::pixel_unshuffle(input, downscale_factor);
17 }
18 } // namespace detail
19 #endif /* DOXYGEN_SHOULD_SKIP_THIS */
20 
21 /// See
22 /// https://pytorch.org/docs/main/nn.functional.html#torch.nn.functional.pixel_shuffle
23 /// about the exact behavior of this functional.
24 ///
25 /// See the documentation for `torch::nn::functional::PixelShuffleFuncOptions`
26 /// class to learn what optional arguments are supported for this functional.
27 ///
28 /// Example:
29 /// ```
30 /// namespace F = torch::nn::functional;
31 /// F::pixel_shuffle(x, F::PixelShuffleFuncOptions(2));
32 /// ```
pixel_shuffle(const Tensor & input,const PixelShuffleFuncOptions & options)33 inline Tensor pixel_shuffle(
34     const Tensor& input,
35     const PixelShuffleFuncOptions& options) {
36   return detail::pixel_shuffle(input, options.upscale_factor());
37 }
38 
pixel_unshuffle(const Tensor & input,const PixelUnshuffleFuncOptions & options)39 inline Tensor pixel_unshuffle(
40     const Tensor& input,
41     const PixelUnshuffleFuncOptions& options) {
42   return detail::pixel_unshuffle(input, options.downscale_factor());
43 }
44 
45 } // namespace functional
46 } // namespace nn
47 } // namespace torch
48