static */
| 1878 | } |
| 1879 | |
| 1880 | /* static */ StatusOr<Shape> ShapeInference::InferFftShape( |
| 1881 | const Shape& in, const FftType fft_type, |
| 1882 | const absl::Span<const int64> fft_length) { |
| 1883 | const int64 fft_rank = fft_length.size(); |
| 1884 | if (fft_rank < 1 || fft_rank > 3) { |
| 1885 | return InvalidArgument("FFT only supports ranks 1-3; got %d.", fft_rank); |
| 1886 | } |
| 1887 | #define RET_CHECK_RANK(x) \ |
| 1888 | if (x.dimensions_size() < fft_rank) { \ |
| 1889 | return InvalidArgument( \ |
| 1890 | "FFT of rank %d requires input of at least " \ |
| 1891 | "same rank; got input of rank %d", \ |
| 1892 | fft_rank, x.dimensions_size()); \ |
| 1893 | } |
| 1894 | switch (fft_type) { |
| 1895 | case FFT: |
| 1896 | case IFFT: |
| 1897 | if (in.element_type() != C64) { |
| 1898 | return InvalidArgument("%s requires complex input type, found %s.", |
| 1899 | FftType_Name(fft_type), |
| 1900 | PrimitiveType_Name(in.element_type())); |
| 1901 | } |
| 1902 | RET_CHECK_RANK(in); |
| 1903 | return in; |
| 1904 | case RFFT: { |
| 1905 | if (in.element_type() != F32) { |
| 1906 | return InvalidArgument("RFFT requires F32 input type, found %s.", |
| 1907 | PrimitiveType_Name(in.element_type())); |
| 1908 | } |
| 1909 | RET_CHECK_RANK(in); |
| 1910 | for (int i = 0; i < fft_rank; i++) { |
| 1911 | if (in.dimensions(in.dimensions_size() - fft_rank + i) != |
| 1912 | fft_length[i]) { |
| 1913 | return InvalidArgument( |
| 1914 | "RFFT requires innermost dimensions match fft_length but " |
| 1915 | "dimension %d is %d and should be %d.", |
| 1916 | in.dimensions_size() - fft_rank + i, |
| 1917 | in.dimensions(in.dimensions_size() - fft_rank + i), |
| 1918 | fft_length[i]); |
| 1919 | } |
| 1920 | } |
| 1921 | Shape result = ShapeUtil::ChangeElementType(in, C64); |
| 1922 | // Preserve the size of zero-sized dimensions. |
| 1923 | if (fft_length[fft_rank - 1] != 0) { |
| 1924 | result.set_dimensions(result.dimensions_size() - 1, |
| 1925 | fft_length[fft_rank - 1] / 2 + 1); |
| 1926 | } |
| 1927 | return result; |
| 1928 | } |
| 1929 | case IRFFT: { |
| 1930 | if (in.element_type() != C64) { |
| 1931 | return InvalidArgument("IRFFT requires C64 input type, found %s.", |
| 1932 | PrimitiveType_Name(in.element_type())); |
| 1933 | } |
| 1934 | RET_CHECK_RANK(in); |
| 1935 | Shape result = ShapeUtil::ComplexComponentShape(in); |
| 1936 | for (int i = 0; i < fft_rank - 1; i++) { |
| 1937 | if (in.dimensions(in.dimensions_size() - fft_rank + i) != |
nothing calls this directly
no test coverage detected