Lines Matching full:dim
162 IntArrayRef dim, int64_t norm, bool onesided) { in fft_r2c_maybe_out() argument
167 return at::_fft_r2c_outf(input, dim, norm, onesided, out_mut); in fft_r2c_maybe_out()
169 return at::_fft_r2c(input, dim, norm, onesided); in fft_r2c_maybe_out()
174 IntArrayRef dim, int64_t norm, SymInt last_dim_size) { in fft_c2r_maybe_out() argument
181 return at::_fft_c2r_symint_outf(input, dim, norm, last_dim_size, out_mut); in fft_c2r_maybe_out()
183 return at::_fft_c2r_symint(input, dim, norm, last_dim_size); in fft_c2r_maybe_out()
188 IntArrayRef dim, int64_t norm, bool forward) { in fft_c2c_maybe_out() argument
193 return at::_fft_c2c_outf(input, dim, norm, forward, out_mut); in fft_c2c_maybe_out()
195 return at::_fft_c2c(input, dim, norm, forward); in fft_c2c_maybe_out()
206 const auto input_dim = input.dim(); in fft_c2r()
207 const auto dim = maybe_wrap_dim(unwrapped_dim, input_dim, /*wrap_scalar=*/false); in fft_c2r() local
208 const auto n = n_opt.value_or(2*(input.sym_sizes()[dim] - 1)); in fft_c2r()
211 input = resize_fft_input(input, dim, n/2 + 1); in fft_c2r()
219 function_name, out, input, dim, static_cast<int64_t>(norm), n); in fft_c2r()
232 const auto input_dim = input.dim(); in fft_r2c()
233 const auto dim = maybe_wrap_dim(unwrapped_dim, input_dim, /*wrap_scalar=*/false); in fft_r2c() local
234 const auto n = n_opt.value_or(input.sym_sizes()[dim]); in fft_r2c()
237 input = resize_fft_input(input, dim, n); in fft_r2c()
244 ret = at::_fft_r2c_out(out, input, dim, static_cast<int64_t>(norm), onesided); in fft_r2c()
246 ret = at::_fft_r2c(input, dim, static_cast<int64_t>(norm), onesided); in fft_r2c()
264 const auto input_dim = input.dim(); in fft_c2c()
265 const auto dim = maybe_wrap_dim(unwrapped_dim, input_dim, /*wrap_scalar=*/false); in fft_c2c() local
266 const auto n = n_opt.value_or(input.sym_sizes()[dim]); in fft_c2c()
269 input = resize_fft_input(input, dim, n); in fft_c2c()
272 return fft_c2c_maybe_out(function_name, out, input, dim, norm, forward); in fft_c2c()
278 DimVector dim; member
281 // Pre-process n-dimensional fft's `s` and `dim` arguments.
285 Tensor input, at::OptionalSymIntArrayRef shape, at::OptionalIntArrayRef dim) { in canonicalize_fft_shape_and_dim_args() argument
286 const int64_t input_dim = input.dim(); in canonicalize_fft_shape_and_dim_args()
290 if (dim) { in canonicalize_fft_shape_and_dim_args()
291 ret.dim.resize(dim->size()); in canonicalize_fft_shape_and_dim_args()
292 std::copy(dim->begin(), dim->end(), ret.dim.begin()); in canonicalize_fft_shape_and_dim_args()
293 maybe_wrap_dims(ret.dim, input_dim, /*wrap_scalars=*/false); in canonicalize_fft_shape_and_dim_args()
296 DimVector copy = ret.dim; in canonicalize_fft_shape_and_dim_args()
303 // Has shape, may have dim in canonicalize_fft_shape_and_dim_args()
304 TORCH_CHECK(!dim || in canonicalize_fft_shape_and_dim_args()
305 dim->size() == shape->size(), in canonicalize_fft_shape_and_dim_args()
306 "When given, dim and shape arguments must have the same length"); in canonicalize_fft_shape_and_dim_args()
312 if (!dim) { in canonicalize_fft_shape_and_dim_args()
313 ret.dim.resize(transform_ndim); in canonicalize_fft_shape_and_dim_args()
314 std::iota(ret.dim.begin(), ret.dim.end(), input_dim - transform_ndim); in canonicalize_fft_shape_and_dim_args()
321 ret.shape[i] = n == -1 ? input_sizes[ret.dim[i]] : n; in canonicalize_fft_shape_and_dim_args()
323 } else if (!dim) { in canonicalize_fft_shape_and_dim_args()
324 // No shape, no dim in canonicalize_fft_shape_and_dim_args()
325 ret.dim.resize(input_dim); in canonicalize_fft_shape_and_dim_args()
326 std::iota(ret.dim.begin(), ret.dim.end(), int64_t{0}); in canonicalize_fft_shape_and_dim_args()
330 // No shape, has dim in canonicalize_fft_shape_and_dim_args()
331 ret.shape.resize(ret.dim.size()); in canonicalize_fft_shape_and_dim_args()
332 for (const auto i : c10::irange(ret.dim.size())) { in canonicalize_fft_shape_and_dim_args()
333 ret.shape[i] = input_sizes[ret.dim[i]]; in canonicalize_fft_shape_and_dim_args()
349 IntArrayRef dim, std::optional<c10::string_view> norm_str, bool forward) { in fftn_c2c() argument
351 Tensor x = resize_fft_input(input, dim, shape); in fftn_c2c()
354 return fft_c2c_maybe_out(fname, out, x, dim, norm, forward); in fftn_c2c()
360 Tensor fft_fft_symint(const Tensor& self, std::optional<SymInt> n, int64_t dim, in fft_fft_symint() argument
363 fft_c2c("fft", {}, self, n, dim, norm, /*forward=*/true) : in fft_fft_symint()
364 fft_r2c("fft", {}, self, n, dim, norm, /*forward=*/true, /*onesided=*/false); in fft_fft_symint()
368 int64_t dim, std::optional<c10::string_view> norm, Tensor& out) { in fft_fft_symint_out() argument
370 fft_c2c("fft", out, self, n, dim, norm, /*forward=*/true); in fft_fft_symint_out()
372 fft_r2c("fft", out, self, n, dim, norm, /*forward=*/true, /*onesided=*/false); in fft_fft_symint_out()
377 Tensor fft_ifft_symint(const Tensor& self, std::optional<SymInt> n, int64_t dim, in fft_ifft_symint() argument
380 fft_c2c("ifft", {}, self, n, dim, norm, /*forward=*/false) : in fft_ifft_symint()
381 fft_r2c("ifft", {}, self, n, dim, norm, /*forward=*/false, /*onesided=*/false); in fft_ifft_symint()
385 int64_t dim, std::optional<c10::string_view> norm, Tensor& out) { in fft_ifft_symint_out() argument
387 fft_c2c("ifft", out, self, n, dim, norm, /*forward=*/false); in fft_ifft_symint_out()
389 fft_r2c("ifft", out, self, n, dim, norm, /*forward=*/false, /*onesided=*/false); in fft_ifft_symint_out()
394 Tensor fft_rfft_symint(const Tensor& self, std::optional<SymInt> n, int64_t dim, in fft_rfft_symint() argument
396 return fft_r2c("rfft", {}, self, n, dim, norm, /*forward=*/true, /*onesided=*/true); in fft_rfft_symint()
400 int64_t dim, std::optional<c10::string_view> norm, Tensor& out) { in fft_rfft_symint_out() argument
401 fft_r2c("rfft", out, self, n, dim, norm, /*forward=*/true, /*onesided=*/true); in fft_rfft_symint_out()
405 Tensor fft_irfft_symint(const Tensor& self, std::optional<SymInt> n, int64_t dim, in fft_irfft_symint() argument
407 return fft_c2r("irfft", {}, self, n, dim, norm, /*forward=*/false); in fft_irfft_symint()
411 int64_t dim, std::optional<c10::string_view> norm, Tensor& out) { in fft_irfft_symint_out() argument
412 fft_c2r("irfft", out, self, n, dim, norm, /*forward=*/false); in fft_irfft_symint_out()
416 Tensor fft_hfft_symint(const Tensor& self, std::optional<SymInt> n, int64_t dim, in fft_hfft_symint() argument
418 return fft_c2r("hfft", {}, self, n, dim, norm, /*forward=*/true); in fft_hfft_symint()
422 int64_t dim, std::optional<c10::string_view> norm, Tensor& out) { in fft_hfft_symint_out() argument
423 fft_c2r("hfft", out, self, n, dim, norm, /*forward=*/true); in fft_hfft_symint_out()
427 Tensor fft_ihfft_symint(const Tensor& self, std::optional<SymInt> n, int64_t dim, in fft_ihfft_symint() argument
429 return fft_r2c("ihfft", {}, self, n, dim, norm, /*forward=*/false, /*onesided=*/true); in fft_ihfft_symint()
433 int64_t dim, std::optional<c10::string_view> norm, Tensor& out) { in fft_ihfft_symint_out() argument
434 fft_r2c("ihfft", out, self, n, dim, norm, /*forward=*/false, /*onesided=*/true); in fft_ihfft_symint_out()
439 at::OptionalIntArrayRef dim, in fft_fftn_symint() argument
441 auto desc = canonicalize_fft_shape_and_dim_args(self, s, dim); in fft_fftn_symint()
444 return fftn_c2c("fftn", {}, input, desc.shape, desc.dim, norm, /*forward=*/true); in fft_fftn_symint()
449 at::OptionalIntArrayRef dim, in fft_fftn_symint_out() argument
451 auto desc = canonicalize_fft_shape_and_dim_args(self, s, dim); in fft_fftn_symint_out()
454 fftn_c2c("fftn", out, input, desc.shape, desc.dim, norm, /*forward=*/true); in fft_fftn_symint_out()
459 at::OptionalIntArrayRef dim, in fft_ifftn_symint() argument
461 auto desc = canonicalize_fft_shape_and_dim_args(self, s, dim); in fft_ifftn_symint()
463 return fftn_c2c("ifftn", {}, input, desc.shape, desc.dim, norm, /*forward=*/false); in fft_ifftn_symint()
468 at::OptionalIntArrayRef dim, in fft_ifftn_symint_out() argument
470 auto desc = canonicalize_fft_shape_and_dim_args(self, s, dim); in fft_ifftn_symint_out()
472 fftn_c2c("ifftn", out, input, desc.shape, desc.dim, norm, /*forward=*/false); in fft_ifftn_symint_out()
478 at::OptionalIntArrayRef dim, in fft_rfftn_impl() argument
481 auto desc = canonicalize_fft_shape_and_dim_args(self, s, dim); in fft_rfftn_impl()
484 Tensor x = resize_fft_input(input, desc.dim, desc.shape); in fft_rfftn_impl()
487 return fft_r2c_maybe_out(fname, out, x, desc.dim, norm, /*onesided=*/true); in fft_rfftn_impl()
491 at::OptionalIntArrayRef dim, in fft_rfftn_symint() argument
493 return fft_rfftn_impl({}, self, s, dim, norm_str); in fft_rfftn_symint()
498 at::OptionalIntArrayRef dim, in fft_rfftn_symint_out() argument
500 fft_rfftn_impl(out, self, s, dim, norm_str); in fft_rfftn_symint_out()
516 const auto last_dim = desc.dim.back(); in canonicalize_fft_c2r_shape_and_dim_args()
530 at::OptionalIntArrayRef dim, in fft_irfftn_impl() argument
534 "irfftn", self, s, dim, last_dim_size); in fft_irfftn_impl()
536 Tensor x = resize_fft_input(input, desc.dim, desc.shape); in fft_irfftn_impl()
539 return fft_c2r_maybe_out(fname, out, x, desc.dim, norm, last_dim_size); in fft_irfftn_impl()
544 at::OptionalIntArrayRef dim, in fft_irfftn_symint() argument
546 return fft_irfftn_impl({}, self, s, dim, norm_str); in fft_irfftn_symint()
551 at::OptionalIntArrayRef dim, in fft_irfftn_symint_out() argument
553 fft_irfftn_impl(out, self, s, dim, norm_str); in fft_irfftn_symint_out()
560 at::OptionalIntArrayRef dim, in fft_hfftn_impl() argument
566 fname, self, s, dim, last_dim_size); in fft_hfftn_impl()
568 auto x = resize_fft_input(input, desc.dim, desc.shape); in fft_hfftn_impl()
573 if (desc.dim.size() > 1) { in fft_hfftn_impl()
574 auto c2c_dims = IntArrayRef(desc.dim).slice(0, desc.dim.size() - 1); in fft_hfftn_impl()
580 const auto last_dim = desc.dim.back(); in fft_hfftn_impl()
588 at::OptionalIntArrayRef dim, in fft_hfftn_symint() argument
590 return fft_hfftn_impl(self, s, dim, norm, {}); in fft_hfftn_symint()
596 at::OptionalIntArrayRef dim, std::optional<c10::string_view> norm, in fft_hfftn_symint_out() argument
598 fft_hfftn_impl(self, s, dim, norm, out); in fft_hfftn_symint_out()
605 const at::OptionalIntArrayRef& dim, in fft_ihfftn_impl() argument
609 auto desc = canonicalize_fft_shape_and_dim_args(self, s, dim); in fft_ihfftn_impl()
612 auto x = resize_fft_input(input, desc.dim, desc.shape); in fft_ihfftn_impl()
616 const auto last_dim = desc.dim.back(); in fft_ihfftn_impl()
618 if (desc.dim.size() == 1) { in fft_ihfftn_impl()
623 auto c2c_dims = IntArrayRef(desc.dim).slice(0, desc.dim.size() - 1); in fft_ihfftn_impl()
630 at::OptionalIntArrayRef dim, in fft_ihfftn_symint() argument
632 return fft_ihfftn_impl(self, s, dim, norm, {}); in fft_ihfftn_symint()
638 at::OptionalIntArrayRef dim, in fft_ihfftn_symint_out() argument
641 fft_ihfftn_impl(self, s, dim, norm, out); in fft_ihfftn_symint_out()
646 IntArrayRef dim, std::optional<c10::string_view> norm) { in fft_fft2_symint() argument
647 return native::fft_fftn_symint(self, s, dim, std::move(norm)); in fft_fft2_symint()
651 IntArrayRef dim, std::optional<c10::string_view> norm, Tensor& out) { in fft_fft2_symint_out() argument
652 return native::fft_fftn_symint_out(self, s, dim, std::move(norm), out); in fft_fft2_symint_out()
656 IntArrayRef dim, std::optional<c10::string_view> norm) { in fft_ifft2_symint() argument
657 return native::fft_ifftn_symint(self, s, dim, std::move(norm)); in fft_ifft2_symint()
661 IntArrayRef dim, std::optional<c10::string_view> norm, Tensor& out) { in fft_ifft2_symint_out() argument
662 return native::fft_ifftn_symint_out(self, s, dim, std::move(norm), out); in fft_ifft2_symint_out()
666 IntArrayRef dim, std::optional<c10::string_view> norm) { in fft_rfft2_symint() argument
667 return native::fft_rfftn_symint(self, s, dim, std::move(norm)); in fft_rfft2_symint()
671 IntArrayRef dim, std::optional<c10::string_view> norm, Tensor& out) { in fft_rfft2_symint_out() argument
672 return native::fft_rfftn_symint_out(self, s, dim, std::move(norm), out); in fft_rfft2_symint_out()
676 IntArrayRef dim, std::optional<c10::string_view> norm) { in fft_irfft2_symint() argument
677 return native::fft_irfftn_symint(self, s, dim, std::move(norm)); in fft_irfft2_symint()
681 IntArrayRef dim, std::optional<c10::string_view> norm, Tensor& out) { in fft_irfft2_symint_out() argument
682 return native::fft_irfftn_symint_out(self, s, dim, std::move(norm), out); in fft_irfft2_symint_out()
686 const Tensor& self, at::OptionalSymIntArrayRef s, IntArrayRef dim, in fft_hfft2_symint_out() argument
688 return native::fft_hfftn_symint_out(self, s, dim, std::move(norm), out); in fft_hfft2_symint_out()
692 IntArrayRef dim, std::optional<c10::string_view> norm) { in fft_hfft2_symint() argument
693 return native::fft_hfftn_symint(self, s, dim, std::move(norm)); in fft_hfft2_symint()
697 const Tensor& self, at::OptionalSymIntArrayRef s, IntArrayRef dim, in fft_ihfft2_symint_out() argument
699 return native::fft_ihfftn_symint_out(self, s, dim, std::move(norm), out); in fft_ihfft2_symint_out()
703 IntArrayRef dim, std::optional<c10::string_view> norm) { in fft_ihfft2_symint() argument
704 return native::fft_ihfftn_symint(self, s, dim, std::move(norm)); in fft_ihfft2_symint()
751 // If an array dim is specified, wraps them according to self.dim().
754 DimVector dim; in default_alldims() local
757 dim.resize(dim_unwrapped.size()); in default_alldims()
758 for (const auto i : c10::irange(dim.size())) { in default_alldims()
759 dim[i] = maybe_wrap_dim(dim_unwrapped[i], self.dim(), /*wrap_scalars=*/false); in default_alldims()
762 dim.resize(self.dim()); in default_alldims()
763 std::iota(dim.begin(), dim.end(), 0); in default_alldims()
765 return dim; in default_alldims()
769 auto dim = default_alldims(x, dim_opt); in fft_fftshift() local
772 SymDimVector shift(dim.size()); in fft_fftshift()
773 for (const auto i : c10::irange(dim.size())) { in fft_fftshift()
774 shift[i] = x_sizes[dim[i]] / 2; in fft_fftshift()
777 return at::roll_symint(x, shift, dim); in fft_fftshift()
781 auto dim = default_alldims(x, dim_opt); in fft_ifftshift() local
784 SymDimVector shift(dim.size()); in fft_ifftshift()
785 for (const auto i : c10::irange(dim.size())) { in fft_ifftshift()
786 shift[i] = (x_sizes[dim[i]] + 1) / 2; in fft_ifftshift()
789 return at::roll_symint(x, shift, dim); in fft_ifftshift()
886 if (self.dim() > 2 || self.dim() < 1) { in stft()
892 if (self.dim() == 1) { in stft()
927 if (window.defined() && (window.dim() != 1 || window.size(0) != win_length)) { in stft()
964 out = at::_fft_c2c(input, input.dim() - 1, static_cast<int64_t>(norm), /*forward=*/true); in stft()
966 out = at::_fft_r2c(input, input.dim() - 1, static_cast<int64_t>(norm), onesided); in stft()
970 if (self.dim() == 1) { in stft()
1056 const auto input_dim = input.dim(); in istft()
1106 if (window.dim() != 1 || window.size(0) != win_length) { in istft()
1130 …input = at::_fft_c2c(input, input.dim() - 1, static_cast<int64_t>(norm), /*forward=*/false); // s… in istft()
1137 …input = at::_fft_c2r(input, input.dim() - 1, static_cast<int64_t>(norm), n_fft); // size: (channe… in istft()
1146 /*dim=*/1, in istft()
1153 /*dim=*/1, in istft()
1202 DimVector dim(dim_.begin(), dim_.end()); in _fft_fill_with_conjugate_symmetry_() local
1203 at::maybe_wrap_dims(dim, input_strides.size(), /*wrap_scalars=*/false); in _fft_fill_with_conjugate_symmetry_()
1205 if (input.numel() == 0 || input_sizes[dim.back()] <= 2) { in _fft_fill_with_conjugate_symmetry_()
1210 dim.erase( in _fft_fill_with_conjugate_symmetry_()
1211 std::remove_if(dim.begin(), dim.end(), [&](int64_t dim) { in _fft_fill_with_conjugate_symmetry_() argument
1212 return (input_sizes[dim] <= 2); in _fft_fill_with_conjugate_symmetry_()
1214 dim.end()); in _fft_fill_with_conjugate_symmetry_()
1222 .declare_static_shape(input_sizes, dim) in _fft_fill_with_conjugate_symmetry_()
1227 const auto ndim = static_cast<int64_t>(iter_strides.size() + dim.size()); in _fft_fill_with_conjugate_symmetry_()
1235 for (const auto i : c10::irange(dim.size())) { in _fft_fill_with_conjugate_symmetry_()
1237 in_strides[iter_strides.size() + i] = input_strides[dim[i]] * element_size; in _fft_fill_with_conjugate_symmetry_()
1238 signal_half_sizes[iter_strides.size() + i] = input_sizes[dim[i]]; in _fft_fill_with_conjugate_symmetry_()
1242 signal_half_sizes.back() = (input_sizes[dim.back()] - 1) / 2; in _fft_fill_with_conjugate_symmetry_()
1247 const auto* in_data = data_ptr + input_strides[dim.back()] * element_size; in _fft_fill_with_conjugate_symmetry_()
1249 input_strides[dim.back()] * (input_sizes[dim.back()] - 1) * element_size); in _fft_fill_with_conjugate_symmetry_()
1274 mirror_dims.reserve(dim.size() - 1); in _fft_fill_with_conjugate_symmetry_()
1277 …dim_permute[i] != ndim - 1) { // Not the last dim, which is mirrored separately with negative str… in _fft_fill_with_conjugate_symmetry_()
1281 TORCH_INTERNAL_ASSERT(mirror_dims.size() == dim.size() - 1); in _fft_fill_with_conjugate_symmetry_()