huggingface/candle
1
0
Fork 0
mirror of https://github.com/huggingface/candle.git synced 2025-06-17 02:58:50 +00:00
Code Issues Packages Projects Releases Wiki Activity

Add the casting operation.

Browse Source
This commit is contained in:
laurent
2023-06-23 21:22:07 +01:00
parent 8ed350dc94
commit 5d44e76e3f
6 changed files with 83 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
src/tensor.rs
View File

@ -602,6 +602,17 @@ impl Tensor {
}
}
pub fn to_dtype(&self, dtype: DType) -> Result<Self> {
let shape = self.shape();
let storage = self.storage.to_dtype(shape, self.stride(), dtype)?;
let op = if self.track_op() {
Some(Op::ToDType(self.clone()))
} else {
None
};
Ok(from_storage(storage, shape.clone(), op, false))
}
pub fn contiguous(&self) -> Result<Tensor> {
if self.is_contiguous() {
Ok(self.clone())
@ -773,6 +784,7 @@ impl Tensor {
}
}
Op::Reshape(node)
| Op::ToDType(node)
| Op::ToDevice(node)
| Op::Transpose(node, _, _)
| Op::Softmax(node, _)
@ -892,6 +904,10 @@ impl Tensor {
*rhs_sum_grad = rhs_sum_grad.add(&rhs_grad)?;
}
Op::Cat(_args, _dim) => return Err(Error::BackwardNotSupported { op: "cat" }),
Op::ToDType(arg) => {
let sum_grad = grads.or_insert(arg)?;
*sum_grad = sum_grad.add(&grad.to_dtype(node.dtype())?)?
}
Op::Affine { arg, mul, .. } => {
let arg_grad = grad.affine(*mul, 0.)?;
let sum_grad = grads.or_insert(arg)?;