Put the spectrum into one big 2D window, not batches of small windows

Author: frutiemax92

src/peft/tuners/fourierft/layer.py

@@ -244,15 +244,15 @@ def update_layer(
         self.fourierft_n_frequency[adapter_name] = n_frequency
         self.fourierft_random_loc_seed[adapter_name] = random_loc_seed
         self.indices[adapter_name] = torch.randperm(
-            self.out_features * self.in_features,
+            self.out_features * self.in_features * self.kW * self.kH,
             generator=torch.Generator().manual_seed(self.fourierft_random_loc_seed[adapter_name]),
         )[:n_frequency]
         self.indices[adapter_name] = torch.stack(
-            [self.indices[adapter_name] // self.in_features, self.indices[adapter_name] % self.in_features], dim=0
+            [self.indices[adapter_name] // (self.in_features * self.kW), self.indices[adapter_name] % (self.in_features * self.kW)], dim=0
         )
         self.fourierft_scaling[adapter_name] = scaling
         # Actual trainable parameters
-        self.fourierft_spectrum[adapter_name] = nn.Parameter(torch.randn(n_frequency, self.kW, self.kH), requires_grad=True)
+        self.fourierft_spectrum[adapter_name] = nn.Parameter(torch.randn(n_frequency), requires_grad=True)
 
         if init_weights:
             self.reset_fourier_parameters(adapter_name)
@@ -310,13 +310,12 @@ def unmerge(self) -> None:
                 self.get_base_layer().weight.data -= self.get_delta_weight(active_adapter)
 
     def get_delta_weight(self, adapter) -> torch.Tensor:
-        # careful: ifft2 does not work with float16 or bfloat16
         spectrum = self.fourierft_spectrum[adapter]
         indices = self.indices[adapter].to(spectrum.device)
-        dense_spectrum = torch.zeros(self.out_features, self.in_features, self.kW, self.kH, device=spectrum.device)
+        dense_spectrum = torch.zeros(self.out_features * self.kH, self.in_features * self.kW, device=spectrum.device)
         dense_spectrum[indices[0, :], indices[1, :]] = spectrum.float()
         delta_weight = torch.fft.ifft2(dense_spectrum).real * self.fourierft_scaling[adapter]
-        return delta_weight
+        return delta_weight.reshape((self.out_features, self.in_features, self.kW, self.kH))
 
     def forward(self, x: torch.Tensor, *args: Any, **kwargs: Any) -> torch.Tensor:
         previous_dtype = x.dtype