Skip to content

Commit

Permalink
PyO3: Add None and Tensor indexing to candle.Tensor (huggingfac…
Browse files Browse the repository at this point in the history 
…e#1098)

* Add proper `None` and `tensor` indexing

* Allow indexing via lists + allow tensor/list indexing outside of first dimension
  • Loading branch information
@LLukas22
LLukas22 authored Oct 20, 2023
1 parent 31ca489 commit b43ab6c
Show file tree
Hide file tree
Showing 2 changed files with 132 additions and 32 deletions.
126 changes: 94 additions & 32 deletions candle-pyo3/src/lib.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -201,6 +201,8 @@ enum Indexer {
Index(usize),
Slice(usize, usize),
Elipsis,
Expand,
IndexSelect(Tensor),
}

#[pymethods]
Expand Down Expand Up @@ -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
Expand All @@ -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
)))
}
}

for (i, item) in tuple.iter().enumerate() {
if item.is_ellipsis() {
// Handle '...' e.g. tensor[..., 0]
if let Ok(tuple) = idx.downcast::<pyo3::types::PyTuple>(py) {
let not_none_count: usize = tuple.iter().filter(|x| !x.is_none()).count();

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"));
}
if not_none_count > dims.len() {
return Err(PyValueError::new_err("provided too many indices"));
}

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();
Expand All @@ -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
}
}
}

Expand Down
38 changes: 38 additions & 0 deletions candle-pyo3/tests/native/test_tensor.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -55,6 +55,7 @@ def test_tensor_can_be_sliced():
assert t[-4:].values() == [5.0, 9.0, 2.0, 6.0]
assert t[:-4].values() == [3.0, 1.0, 4.0, 10.0]
assert t[-4:-2].values() == [5.0, 9.0]
assert t[...].values() == t.values()


def test_tensor_can_be_sliced_2d():
Expand All @@ -76,6 +77,43 @@ def test_tensor_can_be_scliced_3d():
assert t[..., 0:2].values() == [[[1, 2], [5, 6]], [[9, 10], [13, 14]]]


def test_tensor_can_be_expanded_with_none():
t = candle.rand((12, 12))

b = t[None]
assert b.shape == (1, 12, 12)
c = t[:, None, None, :]
assert c.shape == (12, 1, 1, 12)
d = t[None, :, None, :]
assert d.shape == (1, 12, 1, 12)
e = t[None, None, :, :]
assert e.shape == (1, 1, 12, 12)
f = t[:, :, None]
assert f.shape == (12, 12, 1)


def test_tensor_can_be_index_via_tensor():
t = candle.Tensor([[1, 2, 1, 2], [3, 4, 3, 4], [5, 6, 5, 6]])
indexed = t[candle.Tensor([0, 2])]
assert indexed.shape == (2, 4)
assert indexed.values() == [[1, 2, 1, 2], [5, 6, 5, 6]]

indexed = t[:, candle.Tensor([0, 2])]
assert indexed.shape == (3, 2)
assert indexed.values() == [[1, 1], [3, 3], [5, 5]]


def test_tensor_can_be_index_via_list():
t = candle.Tensor([[1, 2, 1, 2], [3, 4, 3, 4], [5, 6, 5, 6]])
indexed = t[[0, 2]]
assert indexed.shape == (2, 4)
assert indexed.values() == [[1, 2, 1, 2], [5, 6, 5, 6]]

indexed = t[:, [0, 2]]
assert indexed.shape == (3, 2)
assert indexed.values() == [[1, 1], [3, 3], [5, 5]]


def test_tensor_can_be_cast_via_to():
t = Tensor(42.0)
assert str(t.dtype) == str(candle.f32)
Expand Down

0 comments on commit b43ab6c

Please sign in to comment.