nato-opt 0.1.0

NATO Optimizer and Spectral Penalties (Rust Port)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85

use tch::Tensor;

pub fn low_pass_filter_gradients(
    named_parameters: &[(String, Tensor)],
    tau_ratio: f64,
    skip_bias: bool,
    in_place: bool,
) {
    for (_, param) in named_parameters.iter() {
        if !param.requires_grad() || !param.grad().defined() {
            continue;
        }

        let grad = param.grad();
        if skip_bias && grad.dim() < 2 {
            continue;
        }

        let shape = grad.size();
        let device = grad.device();
        
        // Convert to complex and FFT
        // Using un-suffixed fftn if complex or standard.
        // Let's use `fft_fftn` if it exists. PyTorch has `torch.fft.fftn`.
        // In tch, it's usually `t.fft_fftn(None, None, None)` where arguments are s, dim, norm.
        // We'll just call `t.fft_fftn::<&[i64], &[i64]>(None, None, None)` - wait, type inference might need help.
        let grad_complex = grad.to_kind(tch::Kind::ComplexFloat);
        
        let grad_fft = grad_complex.fft_fftn(
            None::<&[i64]>, 
            None::<&[i64]>, 
            "backward"
        );
        
        // fftshift
        let mut grad_fft_shifted = grad_fft;
        for (i, &dim) in shape.iter().enumerate() {
            let shift = dim / 2;
            grad_fft_shifted = grad_fft_shifted.roll(&[shift as i64], &[i as i64]);
        }
        
        // Build low-frequency mask
        let mut mask_low = Tensor::ones(shape.as_slice(), (tch::Kind::Bool, device));
        for (axis, &dim) in shape.iter().enumerate() {
            let c = std::cmp::max(1, (dim as f64 * tau_ratio) as i64);
            let start = std::cmp::max(0, dim / 2 - c);
            let end = std::cmp::min(dim, dim / 2 + c);
            
            let mut idx_mask = vec![false; dim as usize];
            for i in start..end {
                idx_mask[i as usize] = true;
            }
            let axis_low = Tensor::from_slice(&idx_mask).to_device(device);
            let mut view_shape = vec![1; shape.len()];
            view_shape[axis] = dim;
            let axis_low = axis_low.view(view_shape.as_slice());
            
            mask_low = mask_low.logical_and(&axis_low);
        }

        let filtered_fft = grad_fft_shifted.where_self(&mask_low, &Tensor::zeros_like(&grad_fft_shifted));

        // ifftshift
        let mut unshifted = filtered_fft;
        for (i, &dim) in shape.iter().enumerate() {
            let shift = (dim + 1) / 2;
            unshifted = unshifted.roll(&[shift as i64], &[i as i64]);
        }

        // ifftn
        let grad_filtered = unshifted.fft_ifftn(
            None::<&[i64]>, 
            None::<&[i64]>, 
            "backward"
        ).real();

        if in_place {
            let _ = param.grad().copy_(&grad_filtered);
        } else {
            // Note: in `tch`, replacing the gradient tensor entirely requires `param.set_grad(grad_filtered)`
            // We'll assume `copy_` is sufficient since usually we want to update the data in place.
            let _ = param.grad().copy_(&grad_filtered);
        }
    }
}