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

Adding size checking when creating a tensor from buffer + shape.

Browse Source
This commit is contained in:
Nicolas Patry
2023-06-22 13:08:57 +02:00
parent a8b6c848e0
commit 449af49b54
2 changed files with 11 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
src/error.rs
View File

@ -12,6 +12,11 @@ pub enum Error {
#[error("the candle crate has not been built with cuda support")]
NotCompiledWithCudaSupport,
#[error(
"Shape mismatch, got buffer of size {buffer_size} which is compatible with shape {shape:?}"
)]
ShapeMismatch { buffer_size: usize, shape: Shape },
#[error("shape mismatch in {op}, lhs: {lhs:?}, rhs: {rhs:?}")]
ShapeMismatchBinaryOp {
lhs: Shape,
@ -40,6 +45,7 @@ pub enum Error {
shape: Shape,
},
// TODO this is temporary when we support arbitrary matmul
#[error("temporary error where matmul doesn't support arbitrary striding")]
UnexpectedStriding,

6
src/tensor.rs
View File

@ -151,7 +151,11 @@ impl Tensor {
device: &Device,
is_variable: bool,
) -> Result<Self> {
// let shape = array.shape()?;
let n: usize = shape.0.iter().product();
let buffer_size: usize = array.shape()?.0.iter().product();
if buffer_size != n {
return Err(Error::ShapeMismatch { buffer_size, shape });
}
let storage = device.storage(array)?;
let stride = shape.stride_contiguous();
let tensor_ = Tensor_ {