Function timestep_embedding

cake-core/src/models/flux/flux1_model.rs:346–361

(t: &Tensor, dim: usize, dtype: DType)

}

pub fn timestep_embedding(t: &Tensor, dim: usize, dtype: DType) -> Result<Tensor> {
    const TIME_FACTOR: f64 = 1000.;
    const MAX_PERIOD: f64 = 10000.;
    if dim % 2 == 1 {
        candle_core::bail!("{dim} is odd")
    }
    let dev = t.device();
    let half = dim / 2;
    let t = (t * TIME_FACTOR)?;
    // Compute frequency vector directly as f32 Vec — avoids arange + to_dtype + mul + exp
    let decay = -MAX_PERIOD.ln() / half as f64;
    let freqs_data: Vec<f32> = (0..half).map(|j| (j as f64 * decay).exp() as f32).collect();
    let freqs = Tensor::new(freqs_data.as_slice(), dev)?.unsqueeze(0)?;
    let args = t.unsqueeze(1)?.to_dtype(DType::F32)?.broadcast_mul(&freqs)?;
    Tensor::cat(&[args.cos()?, args.sin()?], D::Minus1)?.to_dtype(dtype)
}

// ── Positional Embeddings ──────────────────────────────────────────────────────