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

Backprop for narrow.

Browse Source
This commit is contained in:
laurent
2023-06-24 15:17:57 +01:00
parent fbbf3951dd
commit a6ca9baf3c
1 changed files with 9 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
src/tensor.rs
View File

@ -968,9 +968,15 @@ impl Tensor {
let rhs_sum_grad = grads.or_insert(rhs)?;
*rhs_sum_grad = rhs_sum_grad.add(&rhs_grad)?;
}
Op::Cat(_args, _dim) => {
// TODO: Use narrow here.
return Err(Error::BackwardNotSupported { op: "cat" });
Op::Cat(args, dim) => {
let mut start_idx = 0;
for arg in args {
let len = arg.dims()[*dim];
let arg_grad = grad.narrow(*dim, start_idx, len)?;
let sum_grad = grads.or_insert(arg)?;
*sum_grad = sum_grad.add(&arg_grad)?;
start_idx += len;
}
}
Op::ToDType(arg) => {
let sum_grad = grads.or_insert(arg)?;