Group normalization: `(x - mean) / sqrt(var + eps) * weight + bias` per group. Matches candle_nn::GroupNorm::forward() — F32 promotion for F16/BF16. Input: `(batch, channels, ...)`, weight/bias: `(channels,)`.
(
&self,
x: &Tensor,
weight: &Tensor,
bias: &Tensor,
num_groups: usize,
eps: f32,
)
| 329 | /// Matches candle_nn::GroupNorm::forward() — F32 promotion for F16/BF16. |
| 330 | /// Input: `(batch, channels, ...)`, weight/bias: `(channels,)`. |
| 331 | fn group_norm( |
| 332 | &self, |
| 333 | x: &Tensor, |
| 334 | weight: &Tensor, |
| 335 | bias: &Tensor, |
| 336 | num_groups: usize, |
| 337 | eps: f32, |
| 338 | ) -> Result<Tensor> { |
| 339 | use candle_core::DType; |
| 340 | let x_shape = x.dims(); |
| 341 | let x_dtype = x.dtype(); |
| 342 | // Fast path: F32 CPU data — single-pass raw computation |
| 343 | if x_dtype == DType::F32 { |
| 344 | let (b_sz, n_channels) = (x_shape[0], x_shape[1]); |
| 345 | let spatial: usize = x_shape[2..].iter().product(); |
| 346 | let channels_per_group = n_channels / num_groups; |
| 347 | let group_size = channels_per_group * spatial; |
| 348 | let x_data = x.contiguous()?.flatten_all()?.to_vec1::<f32>()?; |
| 349 | let w_data = weight.to_vec1::<f32>()?; |
| 350 | let b_data = bias.to_vec1::<f32>()?; |
| 351 | let mut out = vec![0f32; x_data.len()]; |
| 352 | for batch in 0..b_sz { |
| 353 | let batch_off = batch * n_channels * spatial; |
| 354 | for g in 0..num_groups { |
| 355 | let group_off = batch_off + g * group_size; |
| 356 | let mut sum = 0f64; |
| 357 | let mut sum_sq = 0f64; |
| 358 | for i in 0..group_size { |
| 359 | let v = x_data[group_off + i] as f64; |
| 360 | sum += v; |
| 361 | sum_sq += v * v; |
| 362 | } |
| 363 | let mean = sum / group_size as f64; |
| 364 | let var = sum_sq / group_size as f64 - mean * mean; |
| 365 | let rstd = 1.0 / (var + eps as f64).sqrt(); |
| 366 | for c in 0..channels_per_group { |
| 367 | let ch = g * channels_per_group + c; |
| 368 | let w = w_data[ch] as f64; |
| 369 | let b = b_data[ch] as f64; |
| 370 | let ch_off = group_off + c * spatial; |
| 371 | for s in 0..spatial { |
| 372 | let v = x_data[ch_off + s] as f64; |
| 373 | out[ch_off + s] = ((v - mean) * rstd * w + b) as f32; |
| 374 | } |
| 375 | } |
| 376 | } |
| 377 | } |
| 378 | return Tensor::from_vec(out, x_shape, x.device()); |
| 379 | } |
| 380 | // Fallback: tensor ops with dtype promotion |
| 381 | let (b_sz, n_channels) = (x_shape[0], x_shape[1]); |
| 382 | let hidden_size = x_shape[2..].iter().product::<usize>() * n_channels / num_groups; |
| 383 | let internal_dtype = match x_dtype { |
| 384 | DType::F16 | DType::BF16 => DType::F32, |
| 385 | d => d, |
| 386 | }; |
| 387 | let x = x.reshape((b_sz, num_groups, hidden_size))?; |
| 388 | let x = x.to_dtype(internal_dtype)?; |