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

Support dim indexes in cat.

Browse Source
This commit is contained in:
laurent
2023-07-05 20:39:08 +01:00
parent fc2ffcc72b
commit e2bfbcb79c
2 changed files with 14 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
candle-core/src/tensor.rs
View File

@ -970,10 +970,11 @@ impl Tensor {
self.reshape(dims)
}
pub fn stack<A: AsRef<Tensor>>(args: &[A], dim: usize) -> Result<Self> {
pub fn stack<A: AsRef<Tensor>, D: Dim>(args: &[A], dim: D) -> Result<Self> {
if args.is_empty() {
return Err(Error::OpRequiresAtLeastOneTensor { op: "stack" });
}
let dim = dim.to_index(args[0].as_ref().shape(), "stack")?;
let args = args
.iter()
.map(|t| t.as_ref().unsqueeze(dim))
@ -981,7 +982,7 @@ impl Tensor {
Self::cat(&args, dim)
}
pub fn cat<A: AsRef<Tensor>>(args: &[A], dim: usize) -> Result<Self> {
pub fn cat<A: AsRef<Tensor>, D: Dim>(args: &[A], dim: D) -> Result<Self> {
if args.is_empty() {
return Err(Error::OpRequiresAtLeastOneTensor { op: "cat" });
}
@ -989,6 +990,7 @@ impl Tensor {
if args.len() == 1 {
return Ok(arg0.clone());
}
let dim = dim.to_index(arg0.shape(), "cat")?;
for arg in args {
arg.as_ref().check_dim(dim, "cat")?;
}

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

@ -15,7 +15,7 @@ use anyhow::{Error as E, Result};
use clap::Parser;
use rand::{distributions::Distribution, SeedableRng};
use candle::{DType, Device, Tensor};
use candle::{DType, Device, Tensor, D};
use candle_hub::{api::Api, Repo, RepoType};
use std::collections::HashMap;
use std::sync::{Arc, Mutex};
@ -283,18 +283,18 @@ impl CausalSelfAttention {
dims.push(v / 2);
dims.push(2);
let x = x.reshape(dims)?;
let re_x = x.narrow(candle::D::Minus1, 0, 1)?;
let im_x = x.narrow(candle::D::Minus1, 1, 1)?;
let re_x = x.narrow(D::Minus1, 0, 1)?;
let im_x = x.narrow(D::Minus1, 1, 1)?;
let re_f = freqs_cis
.narrow(candle::D::Minus1, 0, 1)?
.narrow(D::Minus1, 0, 1)?
.broadcast_as(re_x.shape())?;
let im_f = freqs_cis
.narrow(candle::D::Minus1, 1, 1)?
.narrow(D::Minus1, 1, 1)?
.broadcast_as(im_x.shape())?;
let re = ((&re_x * &re_f)? - (&im_x * &im_f)?)?;
let im = ((&re_x * &im_f)? + (&im_x * &re_f)?)?;
let rope = Tensor::cat(&[&re, &im], re.rank() - 1)?;
let rope = rope.flatten_from(candle::D::Minus2)?;
let rope = Tensor::cat(&[&re, &im], D::Minus1)?;
let rope = rope.flatten_from(D::Minus2)?;
Ok(rope)
}
@ -338,7 +338,7 @@ impl CausalSelfAttention {
let att = (q.matmul(&k.t()?)? / (*k_shape.dims().last().unwrap() as f64).sqrt())?;
let mask = self.cache.mask(t)?.broadcast_as(att.shape())?;
let att = masked_fill(&att, &mask, f32::NEG_INFINITY)?;
let att = att.softmax(candle::D::Minus1)?;
let att = att.softmax(D::Minus1)?;
// Convert to contiguous as matmul doesn't support strided vs for now.
let y = att.matmul(&v.contiguous()?)?;
let y = y.transpose(0, 1)?.reshape(&[t, c])?;
@ -424,8 +424,7 @@ fn precompute_freqs_cis(config: &Config, device: &Device) -> Result<Tensor> {
let shape = [1, MAX_SEQ_LEN, n_elem / 2, 1];
let idx_theta_cos = idx_theta.cos()?.reshape(&shape)?;
let idx_theta_sin = idx_theta.sin()?.reshape(&shape)?;
let last_dim = idx_theta_cos.rank() - 1;
Ok(Tensor::cat(&[&idx_theta_cos, &idx_theta_sin], last_dim)?)
Ok(Tensor::cat(&[&idx_theta_cos, &idx_theta_sin], D::Minus1)?)
}
#[derive(Parser, Debug)]
@ -536,7 +535,7 @@ async fn main() -> Result<()> {
let next_token = if let Some(temperature) = args.temperature {
println!("Sampling with temperature {temperature:?}");
let prs = (&logits / temperature)?.softmax(candle::D::Minus1)?;
let prs = (&logits / temperature)?.softmax(D::Minus1)?;
let logits_v: Vec<f32> = prs.to_vec1()?;
let distr = rand::distributions::WeightedIndex::new(&logits_v)?;