src/transformers/models/fnet/modeling_fnet.py/FNetBasicFourierTransform/_init_fourier_transform
class FNetBasicFourierTransform:    def _init_fourier_transform(self, config):
        if not config.use_tpu_fourier_optimizations:
            self.fourier_transform = partial(torch.fft.fftn, dim=(1, 2))
        elif config.max_position_embeddings <= 4096:
            if is_scipy_available():
                self.register_buffer(
                    "dft_mat_hidden", torch.tensor(linalg.dft(config.hidden_size), dtype=torch.complex64)
                )
                self.register_buffer(
                    "dft_mat_seq", torch.tensor(linalg.dft(config.tpu_short_seq_length), dtype=torch.complex64)
                )
                self.fourier_transform = partial(
                    two_dim_matmul, matrix_dim_one=self.dft_mat_seq, matrix_dim_two=self.dft_mat_hidden
                )
            else:
                logging.warning(
                    "SciPy is needed for DFT matrix calculation and is not found. Using TPU optimized fast fourier"
                    " transform instead."
                )
                self.fourier_transform = fftn
        else:
            self.fourier_transform = fftn