| 131 | } |
| 132 | |
| 133 | fn rms_norm_channel(&self, x: &Tensor, weight: &Tensor, eps: f32) -> Result<Tensor> { |
| 134 | // x is (batch, channels, time) — norm over channels at each time step |
| 135 | if x.dtype() == DType::F32 { |
| 136 | let x = x.contiguous()?; |
| 137 | let batch = x.dim(0)?; |
| 138 | let channels = x.dim(1)?; |
| 139 | let time = x.dim(2)?; |
| 140 | let data = x.flatten_all()?.to_vec1::<f32>()?; |
| 141 | let w = weight.to_vec1::<f32>()?; |
| 142 | let eps64 = eps as f64; |
| 143 | let mut out = vec![0f32; data.len()]; |
| 144 | for b in 0..batch { |
| 145 | let batch_off = b * channels * time; |
| 146 | for t in 0..time { |
| 147 | // Gather channel values at this time step |
| 148 | let mut sum_sq = 0f64; |
| 149 | for c in 0..channels { |
| 150 | let v = data[batch_off + c * time + t] as f64; |
| 151 | sum_sq += v * v; |
| 152 | } |
| 153 | let rms = (sum_sq / channels as f64 + eps64).sqrt(); |
| 154 | let inv_rms = 1.0 / rms; |
| 155 | for (c, &wc) in w.iter().enumerate() { |
| 156 | let idx = batch_off + c * time + t; |
| 157 | out[idx] = (data[idx] as f64 * inv_rms * wc as f64) as f32; |
| 158 | } |
| 159 | } |
| 160 | } |
| 161 | return Tensor::from_vec(out, (batch, channels, time), x.device()); |
| 162 | } |
| 163 | // Fallback: transpose approach |
| 164 | x.transpose(1, 2)? |
| 165 | .contiguous() |
| 166 | .and_then(|t| candle_nn::ops::rms_norm(&t, weight, eps))? |
| 167 | .transpose(1, 2) |
| 168 | } |
| 169 | |
| 170 | fn depthwise_conv1d_silu( |
| 171 | &self, |