83 lines
2.6 KiB
Python
83 lines
2.6 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Debug script to compare spectral convolution alone.
|
|
"""
|
|
|
|
import torch
|
|
import numpy as np
|
|
from safetensors.torch import save_file
|
|
|
|
def debug_spectral_conv():
|
|
"""Create a minimal test case for spectral convolution comparison."""
|
|
torch.manual_seed(42)
|
|
|
|
# Simple configuration
|
|
in_ch, out_ch = 2, 2
|
|
modes_h, modes_w = 4, 4
|
|
height, width = 16, 16
|
|
|
|
# Create a simple input
|
|
x = torch.randn(1, in_ch, height, width)
|
|
|
|
# Create weights [in_ch, out_ch, modes_h, modes_w]
|
|
scale = 0.1
|
|
w1_real = torch.randn(in_ch, out_ch, modes_h, modes_w) * scale
|
|
w1_imag = torch.randn(in_ch, out_ch, modes_h, modes_w) * scale
|
|
w2_real = torch.randn(in_ch, out_ch, modes_h, modes_w) * scale
|
|
w2_imag = torch.randn(in_ch, out_ch, modes_h, modes_w) * scale
|
|
|
|
# Step 1: FFT
|
|
x_ft = torch.fft.fft2(x)
|
|
print(f"Input shape: {x.shape}")
|
|
print(f"FFT shape: {x_ft.shape}")
|
|
print(f"FFT[0,0,0,0]: {x_ft[0,0,0,0]}")
|
|
|
|
# Step 2: Create complex weights
|
|
w1 = torch.complex(w1_real, w1_imag)
|
|
w2 = torch.complex(w2_real, w2_imag)
|
|
|
|
# Step 3: Apply weights to frequency domain
|
|
out_ft = torch.zeros_like(x_ft)
|
|
|
|
# Upper portion
|
|
x_upper = x_ft[:, :, :modes_h, :modes_w]
|
|
out_upper = torch.einsum('bihw,iohw->bohw', x_upper, w1)
|
|
out_ft[:, :, :modes_h, :modes_w] = out_upper
|
|
|
|
# Lower portion
|
|
x_lower = x_ft[:, :, -modes_h:, :modes_w]
|
|
out_lower = torch.einsum('bihw,iohw->bohw', x_lower, w2)
|
|
out_ft[:, :, -modes_h:, :modes_w] = out_lower
|
|
|
|
print(f"out_ft[0,0,0,0]: {out_ft[0,0,0,0]}")
|
|
print(f"out_ft[0,0,-1,0]: {out_ft[0,0,-1,0]}")
|
|
|
|
# Step 4: Inverse FFT
|
|
output = torch.fft.ifft2(out_ft).real
|
|
|
|
print(f"\nOutput shape: {output.shape}")
|
|
print(f"Output range: [{output.min():.6f}, {output.max():.6f}]")
|
|
print(f"Output mean: {output.mean():.6f}")
|
|
print(f"Output[0,0,0,:4]: {output[0,0,0,:4]}")
|
|
|
|
# Save for Rust comparison
|
|
save_file({
|
|
"input": x.contiguous(),
|
|
"w1_real": w1_real.contiguous(),
|
|
"w1_imag": w1_imag.contiguous(),
|
|
"w2_real": w2_real.contiguous(),
|
|
"w2_imag": w2_imag.contiguous(),
|
|
"output": output.contiguous(),
|
|
}, "/tmp/spectral_conv_debug.safetensors")
|
|
|
|
print("\nSaved debug data to /tmp/spectral_conv_debug.safetensors")
|
|
|
|
# Also print the raw weight values for first element
|
|
print(f"\nw1[0,0,0,0] = {w1_real[0,0,0,0]:.6f} + {w1_imag[0,0,0,0]:.6f}i")
|
|
print(f"x_ft[0,0,0,0] = {x_ft[0,0,0,0].real:.6f} + {x_ft[0,0,0,0].imag:.6f}i")
|
|
print(f"After einsum out_upper[0,0,0,0] = {out_upper[0,0,0,0].real:.6f} + {out_upper[0,0,0,0].imag:.6f}i")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
debug_spectral_conv()
|