ruda_tensor/ops/modules/
group_normalization.rs

1use super::LayerNormOutput;
2use crate::{Backend, DType, FloatDType, Shape, TensorMetadata, tensor::FloatTensor};
3
4/// Actual logical channel-group geometry, shared by backend dispatch and operation graphs.
5pub struct GroupNormGeometry {
6    /// Original batch extent, including empty batches.
7    pub batch: usize,
8    /// Original channel extent.
9    pub channels: usize,
10    /// Product of all original trailing spatial extents; rank-two input uses one.
11    pub spatial: usize,
12    /// Number of actual elements normalized in each channel group.
13    pub width: usize,
14    /// Number of actual independent batch/group rows.
15    pub rows: usize,
16}
17
18/// Validate actual `[batch, channels, ...]` input and evenly divided nonempty groups.
19pub fn geometry(shape: &Shape, groups: usize) -> GroupNormGeometry {
20    assert!(shape.num_dims() >= 2, "GroupNorm requires batch and channel axes");
21    let channels = shape[1];
22    assert!(groups > 0 && channels > 0 && channels % groups == 0, "GroupNorm channels must divide into positive groups");
23    let spatial = shape[2..].iter().try_fold(1usize, |size, &extent| size.checked_mul(extent)).expect("GroupNorm spatial overflow");
24    let width = (channels / groups).checked_mul(spatial).expect("GroupNorm group width overflow");
25    assert!(width > 0, "GroupNorm normalized groups must be nonempty");
26    let rows = shape[0].checked_mul(groups).expect("GroupNorm row count overflow");
27    rows.checked_mul(width).expect("GroupNorm element count overflow");
28    GroupNormGeometry { batch: shape[0], channels, spatial, width, rows }
29}
30
31fn affine<B: Backend>(value: &FloatTensor<B>, channels: usize) {
32    assert_eq!(value.shape(), Shape::new([channels]), "GroupNorm affine must match actual channels");
33}
34
35/// Same-device biased-variance GroupNorm with FP32/FP64 saved `[batch, groups]` statistics.
36pub fn group_norm_with_stats<B: Backend>(input: FloatTensor<B>, gamma: Option<FloatTensor<B>>, beta: Option<FloatTensor<B>>,
37    groups: usize, epsilon: f64) -> LayerNormOutput<B> {
38    let shape = input.shape();
39    let info = geometry(&shape, groups);
40    for value in gamma.iter().chain(beta.iter()) { affine::<B>(value, info.channels); }
41    let storage: FloatDType = input.dtype().into();
42    let compute = if input.dtype() == DType::F64 { FloatDType::F64 } else { FloatDType::F32 };
43    if info.rows == 0 {
44        let device = B::float_device(&input);
45        return LayerNormOutput { output: input, mean: B::float_zeros(Shape::new([0, groups]), &device, compute),
46            rstd: B::float_zeros(Shape::new([0, groups]), &device, compute) };
47    }
48    let input = B::float_reshape(B::float_cast(input, compute), Shape::new([info.rows, info.width]));
49    let mean = B::float_div_scalar(B::float_sum_dim(input.clone(), 1), (info.width as f64).into());
50    let centered = B::float_sub(input, mean.clone());
51    let variance = B::float_div_scalar(B::float_sum_dim(B::float_mul(centered.clone(), centered.clone()), 1), (info.width as f64).into());
52    let rstd = B::float_recip(B::float_sqrt(B::float_add_scalar(variance, epsilon.into())));
53    let mut output = B::float_reshape(B::float_mul(centered, rstd.clone()), Shape::new([info.batch, info.channels, info.spatial]));
54    if let Some(gamma) = gamma {
55        output = B::float_mul(output, B::float_reshape(B::float_cast(gamma, compute), Shape::new([1, info.channels, 1])));
56    }
57    if let Some(beta) = beta {
58        output = B::float_add(output, B::float_reshape(B::float_cast(beta, compute), Shape::new([1, info.channels, 1])));
59    }
60    LayerNormOutput { output: B::float_reshape(B::float_cast(output, storage), shape),
61        mean: B::float_reshape(mean, Shape::new([info.batch, groups])), rstd: B::float_reshape(rstd, Shape::new([info.batch, groups])) }
62}
63
64/// Independently requested input/actual-weight/bias derivatives on the same backend.
65pub fn group_norm_backward_select<B: Backend>(input: FloatTensor<B>, gamma: Option<FloatTensor<B>>, grad: FloatTensor<B>,
66    mean: FloatTensor<B>, rstd: FloatTensor<B>, groups: usize, mask: [bool; 3]) -> [Option<FloatTensor<B>>; 3] {
67    if mask == [false; 3] { return [None, None, None]; }
68    let shape = input.shape();
69    let info = geometry(&shape, groups);
70    if let Some(gamma) = &gamma { affine::<B>(gamma, info.channels); }
71    assert!(!mask[1] || gamma.is_some(), "GroupNorm weight gradient requires an actual weight");
72    assert_eq!(grad.shape(), shape, "GroupNorm gradient shape differs");
73    assert_eq!(mean.shape(), Shape::new([info.batch, groups]), "GroupNorm mean shape differs");
74    assert_eq!(rstd.shape(), mean.shape(), "GroupNorm reciprocal deviation shape differs");
75    let storage: FloatDType = input.dtype().into();
76    let forward_compute = if input.dtype() == DType::F64 { FloatDType::F64 } else { FloatDType::F32 };
77    let compute = if [&input, &grad, &mean, &rstd].into_iter().chain(gamma.iter()).any(|value| value.dtype() == DType::F64) {
78        FloatDType::F64
79    } else { FloatDType::F32 };
80    if info.rows == 0 {
81        let device = B::float_device(&input);
82        return [mask[0].then_some(input), mask[1].then(|| B::float_zeros(Shape::new([info.channels]), &device,
83            gamma.as_ref().expect("requested GroupNorm weight").dtype().into())),
84            mask[2].then(|| B::float_zeros(Shape::new([info.channels]), &device, compute))];
85    }
86    let grad = B::float_reshape(B::float_cast(grad, compute), Shape::new([info.batch, info.channels, info.spatial]));
87    let bias_grad = mask[2].then(|| B::float_reshape(B::float_sum_dim(B::float_sum_dim(grad.clone(), 0), 2), Shape::new([info.channels])));
88    if !mask[0] && !mask[1] { return [None, None, bias_grad]; }
89    let input = B::float_reshape(B::float_cast(input, forward_compute), Shape::new([info.rows, info.width]));
90    let mean = B::float_reshape(B::float_cast(mean, forward_compute), Shape::new([info.rows, 1]));
91    let rstd = B::float_reshape(B::float_cast(rstd, forward_compute), Shape::new([info.rows, 1]));
92    let normalized = B::float_cast(B::float_mul(B::float_sub(input, mean), rstd.clone()), compute);
93    let normalized_channels = B::float_reshape(normalized.clone(), Shape::new([info.batch, info.channels, info.spatial]));
94    let input_grad = mask[0].then(|| {
95        let scaled = if let Some(gamma) = &gamma {
96            let gamma = B::float_cast(B::float_cast(gamma.clone(), forward_compute), compute);
97            B::float_mul(grad.clone(), B::float_reshape(gamma, Shape::new([1, info.channels, 1])))
98        } else { grad.clone() };
99        let scaled = B::float_reshape(scaled, Shape::new([info.rows, info.width]));
100        let average = B::float_div_scalar(B::float_sum_dim(scaled.clone(), 1), (info.width as f64).into());
101        let product = B::float_div_scalar(B::float_sum_dim(B::float_mul(scaled.clone(), normalized.clone()), 1), (info.width as f64).into());
102        let output = B::float_mul(B::float_cast(rstd, compute), B::float_sub(B::float_sub(scaled, average), B::float_mul(normalized, product)));
103        B::float_reshape(B::float_cast(output, storage), shape)
104    });
105    let weight_grad = mask[1].then(|| {
106        let output = B::float_sum_dim(B::float_sum_dim(B::float_mul(grad, normalized_channels), 0), 2);
107        B::float_reshape(B::float_cast(output, gamma.as_ref().expect("requested GroupNorm weight").dtype().into()), Shape::new([info.channels]))
108    });
109    [input_grad, weight_grad, bias_grad]
110}
ruda_tensor/ops/modules/group_normalization.rs

ruda_tensor/ops/modules/
group_normalization.rs
