huggingface/candle
1
0
Fork 0
mirror of https://github.com/huggingface/candle.git synced 2025-06-16 10:38:54 +00:00
Code Issues Packages Projects Releases Wiki Activity

Use arange in the examples. (#146)

Browse Source
This commit is contained in:
Laurent Mazare
2023-07-12 12:12:34 +01:00
committed by GitHub
parent 20599172ac
commit eae646d322
3 changed files with 7 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
candle-examples/examples/falcon/model.rs
View File

@ -166,8 +166,7 @@ impl FalconRotaryEmbedding {
}
_ => {}
}
let t: Vec<_> = (0..seq_len).map(|c| c as u32).collect();
let t = Tensor::new(t.as_slice(), device)?.to_dtype(dtype)?;
let t = Tensor::arange(0, seq_len as u32, device)?.to_dtype(dtype)?;
let inv_freq = self.inv_freq.to_dtype(dtype)?;
let freqs = t.unsqueeze(1)?.matmul(&inv_freq.unsqueeze(0)?)?;
let emb = Tensor::cat(&[&freqs, &freqs], D::Minus1)?;

7
candle-examples/examples/llama/main.rs
View File

@ -86,11 +86,10 @@ fn precompute_freqs_cis(config: &Config, device: &Device) -> Result<Tensor> {
.step_by(2)
.map(|i| 1f32 / 10000f32.powf(i as f32 / n_elem as f32))
.collect();
let arange: Vec<_> = (0..MAX_SEQ_LEN).map(|c| c as f32).collect();
let theta = Tensor::new(theta.as_slice(), device)?;
let arange = Tensor::new(arange.as_slice(), device)?;
let idx_theta = arange
.reshape((arange.elem_count(), 1))?
let idx_theta = Tensor::arange(0, MAX_SEQ_LEN as u32, device)?
.to_dtype(DType::F32)?
.reshape((MAX_SEQ_LEN, 1))?
.matmul(&theta.reshape((1, theta.elem_count()))?)?;
let shape = [1, MAX_SEQ_LEN, n_elem / 2, 1];
let idx_theta_cos = idx_theta.cos()?.reshape(&shape)?;

5
candle-examples/examples/whisper/model.rs
View File

@ -217,9 +217,10 @@ fn sinusoids(length: usize, channels: usize) -> Result<Tensor> {
let inv_timescales: Vec<_> = (0..channels / 2)
.map(|i| (i as f32 * (-log_timescale_increment)).exp())
.collect();
let arange: Vec<_> = (0..length).map(|c| c as f32).collect();
let inv_timescales = Tensor::new(inv_timescales.as_slice(), &Device::Cpu)?.unsqueeze(0)?;
let arange = Tensor::new(arange.as_slice(), &Device::Cpu)?.unsqueeze(1)?;
let arange = Tensor::arange(0, length as u32, &Device::Cpu)?
.to_dtype(candle::DType::F32)?
.unsqueeze(1)?;
let sh = (length, channels / 2);
let scaled_time = (arange.broadcast_as(sh)? * inv_timescales.broadcast_as(sh)?)?;
let sincos = Tensor::cat(&[scaled_time.sin()?, scaled_time.cos()?], 1)?;