Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

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

Merged
merged 3 commits into from
Oct 20, 2023
Merged

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
caf42e9
Add proper `None` and `tensor` indexing
LLukas22 Oct 15, 2023
c161312
Merge remote-tracking branch 'upstream/main' into feat/better-tensor-…
LLukas22 Oct 20, 2023
50fba7f
Allow indexing via lists + allow tensor/list indexing outside of firs…
LLukas22 Oct 20, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
126 changes: 94 additions & 32 deletions candle-pyo3/src/lib.rs
View file
Open in desktop
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
Open in desktop
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
Loading