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

PyO3: Add None and Tensor indexing to candle.Tensor (#1098)

Browse Source 
* Add proper `None` and `tensor` indexing

* Allow indexing via lists + allow tensor/list indexing outside of first dimension
This commit is contained in:
Lukas Kreussel
2023-10-20 10:59:00 +02:00
committed by GitHub
parent 31ca4897bb
commit b43ab6cd1d
2 changed files with 132 additions and 32 deletions
Download Patch File Download Diff File Expand all files Collapse all files

126
candle-pyo3/src/lib.rs
View File

@ -201,6 +201,8 @@ enum Indexer {
Index(usize),
Slice(usize, usize),
Elipsis,
Expand,
IndexSelect(Tensor),
}
#[pymethods]
@ -450,7 +452,7 @@ impl PyTensor {
let mut indexers: Vec<Indexer> = vec![];
let dims = self.0.shape().dims();
let to_absolute_index = |index: isize, current_dim: usize| {
fn to_absolute_index(index: isize, current_dim: usize, dims: &[usize]) -> PyResult<usize> {
// Convert a relative index to an absolute index e.g. tensor[-1] -> tensor[0]
let actual_index = if index < 0 {
dims[current_dim] as isize + index
@ -460,48 +462,92 @@ impl PyTensor {
// Check that the index is in range
if actual_index < 0 || actual_index >= dims[current_dim] as isize {
return Err(PyTypeError::new_err(format!(
return Err(PyValueError::new_err(format!(
"index out of range for dimension '{i}' with indexer '{value}'",
i = current_dim,
value = index
)));
}
Ok(actual_index as usize)
};
if let Ok(index) = idx.extract(py) {
// Handle a single index e.g. tensor[0] or tensor[-1]
indexers.push(Indexer::Index(to_absolute_index(index, 0)?));
} else if let Ok(slice) = idx.downcast::<pyo3::types::PySlice>(py) {
// Handle a single slice e.g. tensor[0:1] or tensor[0:-1]
let index = slice.indices(dims[0] as c_long)?;
indexers.push(Indexer::Slice(index.start as usize, index.stop as usize));
} else if let Ok(tuple) = idx.downcast::<pyo3::types::PyTuple>(py) {
// Handle multiple indices e.g. tensor[0,0] or tensor[0:1,0:1]
}
if tuple.len() > dims.len() {
return Err(PyTypeError::new_err("provided too many indices"));
fn extract_indexer(
py_indexer: &PyAny,
current_dim: usize,
dims: &[usize],
index_argument_count: usize,
) -> PyResult<(Indexer, usize)> {
if let Ok(index) = py_indexer.extract() {
// Handle a single index e.g. tensor[0] or tensor[-1]
Ok((
Indexer::Index(to_absolute_index(index, current_dim, dims)?),
current_dim + 1,
))
} else if let Ok(slice) = py_indexer.downcast::<pyo3::types::PySlice>() {
// Handle a single slice e.g. tensor[0:1] or tensor[0:-1]
let index = slice.indices(dims[current_dim] as c_long)?;
Ok((
Indexer::Slice(index.start as usize, index.stop as usize),
current_dim + 1,
))
} else if let Ok(tensor) = py_indexer.extract::<PyTensor>() {
// Handle a tensor as indices e.g. tensor[tensor([0,1])]
let t = tensor.0;
if t.rank() != 1 {
return Err(PyTypeError::new_err(
"multi-dimensional tensor indexing is not supported",
));
}
Ok((Indexer::IndexSelect(t), current_dim + 1))
} else if let Ok(list) = py_indexer.downcast::<pyo3::types::PyList>() {
// Handle a list of indices e.g. tensor[[0,1]]
let mut indexes = vec![];
for item in list.iter() {
let index = item.extract::<i64>()?;
indexes.push(index);
}
Ok((
Indexer::IndexSelect(
Tensor::from_vec(indexes, list.len(), &Device::Cpu).map_err(wrap_err)?,
),
current_dim + 1,
))
} else if py_indexer.is_ellipsis() {
// Handle '...' e.g. tensor[..., 0]
if current_dim > 0 {
return Err(PyTypeError::new_err(
"Ellipsis ('...') can only be used at the start of an indexing operation",
));
}
Ok((Indexer::Elipsis, dims.len() - (index_argument_count - 1)))
} else if py_indexer.is_none() {
// Handle None e.g. tensor[None, 0]
Ok((Indexer::Expand, current_dim))
} else {
Err(PyTypeError::new_err(format!(
"unsupported indexer {}",
py_indexer
)))
}
}
if let Ok(tuple) = idx.downcast::<pyo3::types::PyTuple>(py) {
let not_none_count: usize = tuple.iter().filter(|x| !x.is_none()).count();
if not_none_count > dims.len() {
return Err(PyValueError::new_err("provided too many indices"));
}
for (i, item) in tuple.iter().enumerate() {
if item.is_ellipsis() {
// Handle '...' e.g. tensor[..., 0]
if i > 0 {
return Err(PyTypeError::new_err("Ellipsis ('...') can only be used at the start of an indexing operation"));
}
indexers.push(Indexer::Elipsis);
} else if let Ok(slice) = item.downcast::<pyo3::types::PySlice>() {
// Handle slice
let index = slice.indices(dims[i] as c_long)?;
indexers.push(Indexer::Slice(index.start as usize, index.stop as usize));
} else if let Ok(index) = item.extract::<isize>() {
indexers.push(Indexer::Index(to_absolute_index(index, i)?));
} else {
return Err(PyTypeError::new_err("unsupported index"));
}
let mut current_dim = 0;
for item in tuple.iter() {
let (indexer, new_current_dim) =
extract_indexer(item, current_dim, dims, not_none_count)?;
current_dim = new_current_dim;
indexers.push(indexer);
}
} else {
return Err(PyTypeError::new_err("unsupported index"));
let (indexer, _) = extract_indexer(idx.downcast::<PyAny>(py)?, 0, dims, 1)?;
indexers.push(indexer);
}
let mut x = self.0.clone();
@ -526,6 +572,22 @@ impl PyTensor {
current_dim += dims.len() - (indexers.len() - 1);
x
}
Indexer::Expand => {
// Expand is a special case, it means that a new dimension should be added => unsqueeze and advance the current_dim
let out = x.unsqueeze(current_dim).map_err(wrap_err)?;
current_dim += 1;
out
}
Indexer::IndexSelect(indexes) => {
let out = x
.index_select(
&indexes.to_device(x.device()).map_err(wrap_err)?,
current_dim,
)
.map_err(wrap_err)?;
current_dim += 1;
out
}
}
}