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

onnx: implement Size op (#2316)

Browse Source
This commit is contained in:
shua
2024-07-07 19:56:36 +02:00
committed by GitHub
parent 9182f9f5c2
commit eec11ce2ce
2 changed files with 52 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
candle-onnx/src/eval.rs
View File

@ -643,6 +643,13 @@ fn simple_eval_(
let dims = Tensor::from_vec(dims, xs.rank(), xs.device())?;
values.insert(node.output[0].clone(), dims);
}
// https://github.com/onnx/onnx/blob/main/docs/Operators.md#Size
"Size" => {
let data = get(&node.input[0])?;
let size: usize = data.dims().iter().product();
let output = Tensor::from_slice(&[size as i64], (), data.device())?;
values.insert(node.output[0].clone(), output);
}
// https://github.com/onnx/onnx/blob/main/docs/Operators.md#Sqrt
"Sqrt" => {
let xs = get(&node.input[0])?;