huggingface/candle
1
0
Fork 0
mirror of https://github.com/huggingface/candle.git synced 2025-06-18 19:47:12 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #26 from LaurentMazare/narrow-grad

Browse Source 
Add the grad for narrow.
This commit is contained in:
Laurent Mazare
2023-06-28 11:46:13 +01:00
committed by GitHub
parent 2998ff6ef7 666d6dbcac
commit d461d9d751
1 changed files with 25 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

27
candle-core/src/backprop.rs
View File

@ -208,8 +208,31 @@ impl Tensor {
let sum_grad = grads.or_insert(arg)?; let sum_grad = grads.or_insert(arg)?;
*sum_grad = sum_grad.sub(&grad)? *sum_grad = sum_grad.sub(&grad)?
} }
Op::Narrow(_arg, _, _, _) => { &Op::Narrow(ref arg, dim, start_idx, len) => {
return Err(Error::BackwardNotSupported { op: "narrow" }) let arg_dims = arg.dims();
let left_pad = if start_idx == 0 {
None
} else {
let mut dims = arg_dims.to_vec();
dims[dim] = start_idx;
Some(Tensor::zeros(dims, grad.dtype(), &grad.device())?)
};
let right_pad = arg_dims[dim] - start_idx - len;
let right_pad = if right_pad == 0 {
None
} else {
let mut dims = arg_dims.to_vec();
dims[dim] = right_pad;
Some(Tensor::zeros(dims, grad.dtype(), &grad.device())?)
};
let arg_grad = match (left_pad, right_pad) {
(None, None) => grad,
(Some(l), None) => Tensor::cat(&[&l, &grad], dim)?,
(None, Some(r)) => Tensor::cat(&[&grad, &r], dim)?,
(Some(l), Some(r)) => Tensor::cat(&[&l, &grad, &r], dim)?,
};
let sum_grad = grads.or_insert(arg)?;
*sum_grad = sum_grad.add(&arg_grad)?
} }
Op::Softmax(_arg, _) => { Op::Softmax(_arg, _) => {
return Err(Error::BackwardNotSupported { op: "softmax" }) return Err(Error::BackwardNotSupported { op: "softmax" })