diff --git a/candle-examples/Cargo.toml b/candle-examples/Cargo.toml index 7372e24f20..b1913541fb 100644 --- a/candle-examples/Cargo.toml +++ b/candle-examples/Cargo.toml @@ -21,7 +21,7 @@ half = { workspace = true, optional = true } image = { workspace = true } intel-mkl-src = { workspace = true, optional = true } num-traits = { workspace = true } -pyo3 = { version = "0.19.0", features = ["auto-initialize"], optional = true } +pyo3 = { version = "0.20.0", features = ["auto-initialize"], optional = true } rayon = { workspace = true } safetensors = { workspace = true } serde = { workspace = true } diff --git a/candle-pyo3/Cargo.toml b/candle-pyo3/Cargo.toml index 488404bf06..b04524040e 100644 --- a/candle-pyo3/Cargo.toml +++ b/candle-pyo3/Cargo.toml @@ -19,10 +19,10 @@ candle = { path = "../candle-core", version = "0.3.0", package = "candle-core" } candle-nn = { path = "../candle-nn", version = "0.3.0" } half = { workspace = true } intel-mkl-src = { workspace = true, optional = true } -pyo3 = { version = "0.19.0", features = ["extension-module", "abi3-py38"] } +pyo3 = { version = "0.20.0", features = ["extension-module", "abi3-py38"] } [build-dependencies] -pyo3-build-config = "0.19" +pyo3-build-config = "0.20" [features] default = [] diff --git a/candle-pyo3/py_src/candle/__init__.pyi b/candle-pyo3/py_src/candle/__init__.pyi index 437221683b..35b17680b9 100644 --- a/candle-pyo3/py_src/candle/__init__.pyi +++ b/candle-pyo3/py_src/candle/__init__.pyi @@ -1,7 +1,7 @@ # Generated content DO NOT EDIT from typing import Any, Callable, Dict, List, Optional, Tuple, Union, Sequence from os import PathLike -from candle.typing import _ArrayLike, Device, Scalar, Index +from candle.typing import _ArrayLike, Device, Scalar, Index, Shape class bf16(DType): pass @@ -26,21 +26,21 @@ class i64(DType): pass @staticmethod -def ones(shape: Sequence[int], dtype: Optional[DType] = None, device: Optional[Device] = None) -> Tensor: +def ones(*shape: Shape, dtype: Optional[DType] = None, device: Optional[Device] = None) -> Tensor: """ Creates a new tensor filled with ones. """ pass @staticmethod -def rand(shape: Sequence[int], device: Optional[Device] = None) -> Tensor: +def rand(*shape: Shape, device: Optional[Device] = None) -> Tensor: """ Creates a new tensor with random values. """ pass @staticmethod -def randn(shape: Sequence[int], device: Optional[Device] = None) -> Tensor: +def randn(*shape: Shape, device: Optional[Device] = None) -> Tensor: """ Creates a new tensor with random values from a normal distribution. """ @@ -67,7 +67,7 @@ class u8(DType): pass @staticmethod -def zeros(shape: Sequence[int], dtype: Optional[DType] = None, device: Optional[Device] = None) -> Tensor: +def zeros(*shape: Shape, dtype: Optional[DType] = None, device: Optional[Device] = None) -> Tensor: """ Creates a new tensor filled with zeros. """ @@ -174,7 +174,7 @@ class Tensor: Adds the two tensors, while broadcasting the right-hand-side tensor to match the shape of the left-hand-side tensor. """ pass - def broadcast_as(self, shape: Sequence[int]) -> Tensor: + def broadcast_as(self, *shape: Shape) -> Tensor: """ Broadcasts the tensor to the given shape. """ @@ -184,7 +184,7 @@ class Tensor: Divides the two tensors, while broadcasting the right-hand-side tensor to match the shape of the left-hand-side tensor. """ pass - def broadcast_left(self, shape: Sequence[int]) -> Tensor: + def broadcast_left(self, *shape: Shape) -> Tensor: """ Broadcasts the tensor to the given shape, adding new dimensions on the left. """ @@ -329,7 +329,7 @@ class Tensor: Get the `recip` of the tensor. """ pass - def reshape(self, shape: Sequence[int]) -> Tensor: + def reshape(self, *shape: Shape) -> Tensor: """ Reshapes the tensor to the given shape. """ diff --git a/candle-pyo3/py_src/candle/functional/__init__.pyi b/candle-pyo3/py_src/candle/functional/__init__.pyi index 5bf5c4c31f..4f7c2aa65a 100644 --- a/candle-pyo3/py_src/candle/functional/__init__.pyi +++ b/candle-pyo3/py_src/candle/functional/__init__.pyi @@ -1,7 +1,7 @@ # Generated content DO NOT EDIT from typing import Any, Callable, Dict, List, Optional, Tuple, Union, Sequence from os import PathLike -from candle.typing import _ArrayLike, Device, Scalar, Index +from candle.typing import _ArrayLike, Device, Scalar, Index, Shape from candle import Tensor, DType, QTensor @staticmethod diff --git a/candle-pyo3/py_src/candle/typing/__init__.py b/candle-pyo3/py_src/candle/typing/__init__.py index 66bc3d8aba..b2262a97d0 100644 --- a/candle-pyo3/py_src/candle/typing/__init__.py +++ b/candle-pyo3/py_src/candle/typing/__init__.py @@ -18,3 +18,5 @@ Scalar = Union[int, float] Index = Union[int, slice, None, "Ellipsis"] + +Shape = Union[int, Sequence[int]] diff --git a/candle-pyo3/py_src/candle/utils/__init__.pyi b/candle-pyo3/py_src/candle/utils/__init__.pyi index d3b9376675..4ee51c290b 100644 --- a/candle-pyo3/py_src/candle/utils/__init__.pyi +++ b/candle-pyo3/py_src/candle/utils/__init__.pyi @@ -1,7 +1,7 @@ # Generated content DO NOT EDIT from typing import Any, Callable, Dict, List, Optional, Tuple, Union, Sequence from os import PathLike -from candle.typing import _ArrayLike, Device, Scalar, Index +from candle.typing import _ArrayLike, Device, Scalar, Index, Shape from candle import Tensor, DType, QTensor @staticmethod diff --git a/candle-pyo3/src/lib.rs b/candle-pyo3/src/lib.rs index 6d4de80bfb..41c4577fee 100644 --- a/candle-pyo3/src/lib.rs +++ b/candle-pyo3/src/lib.rs @@ -16,26 +16,13 @@ extern crate accelerate_src; use ::candle::{quantized::QTensor, DType, Device, Tensor, WithDType}; +mod shape; +use shape::{PyShape, PyShapeWithHole}; + pub fn wrap_err(err: ::candle::Error) -> PyErr { PyErr::new::(format!("{err:?}")) } -#[derive(Clone, Debug)] -struct PyShape(Vec); - -impl<'source> pyo3::FromPyObject<'source> for PyShape { - fn extract(ob: &'source PyAny) -> PyResult { - let dims: Vec = pyo3::FromPyObject::extract(ob)?; - Ok(PyShape(dims)) - } -} - -impl From for ::candle::Shape { - fn from(val: PyShape) -> Self { - val.0.into() - } -} - #[derive(Clone, Debug)] #[pyclass(name = "Tensor")] /// A `candle` tensor. @@ -684,25 +671,37 @@ impl PyTensor { Ok(Self(tensor)) } - #[pyo3(text_signature = "(self, shape:Sequence[int])")] + #[pyo3(signature=(*shape), text_signature = "(self, *shape:Shape)")] /// Reshapes the tensor to the given shape. /// &RETURNS&: Tensor - fn reshape(&self, shape: PyShape) -> PyResult { - Ok(PyTensor(self.0.reshape(shape).map_err(wrap_err)?)) + fn reshape(&self, shape: PyShapeWithHole) -> PyResult { + Ok(PyTensor( + self.0 + .reshape(shape.to_absolute(&self.0)?) + .map_err(wrap_err)?, + )) } - #[pyo3(text_signature = "(self, shape:Sequence[int])")] + #[pyo3(signature=(*shape), text_signature = "(self, *shape:Shape)")] /// Broadcasts the tensor to the given shape. /// &RETURNS&: Tensor - fn broadcast_as(&self, shape: PyShape) -> PyResult { - Ok(PyTensor(self.0.broadcast_as(shape).map_err(wrap_err)?)) + fn broadcast_as(&self, shape: PyShapeWithHole) -> PyResult { + Ok(PyTensor( + self.0 + .broadcast_as(shape.to_absolute(&self.0)?) + .map_err(wrap_err)?, + )) } - #[pyo3(text_signature = "(self, shape:Sequence[int])")] + #[pyo3(signature=(*shape), text_signature = "(self, *shape:Shape)")] /// Broadcasts the tensor to the given shape, adding new dimensions on the left. /// &RETURNS&: Tensor - fn broadcast_left(&self, shape: PyShape) -> PyResult { - Ok(PyTensor(self.0.broadcast_left(shape).map_err(wrap_err)?)) + fn broadcast_left(&self, shape: PyShapeWithHole) -> PyResult { + Ok(PyTensor( + self.0 + .broadcast_left(shape.to_absolute(&self.0)?) + .map_err(wrap_err)?, + )) } #[pyo3(text_signature = "(self, dim:int)")] @@ -915,21 +914,21 @@ impl PyTensor { } if let Some(kwargs) = kwargs { - if let Some(any) = kwargs.get_item("dtype") { + if let Ok(Some(any)) = kwargs.get_item("dtype") { handle_duplicates( &mut dtype, any.extract::(), "cannot specify multiple dtypes", )?; } - if let Some(any) = kwargs.get_item("device") { + if let Ok(Some(any)) = kwargs.get_item("device") { handle_duplicates( &mut device, any.extract::(), "cannot specify multiple devices", )?; } - if let Some(any) = kwargs.get_item("other") { + if let Ok(Some(any)) = kwargs.get_item("other") { handle_duplicates( &mut other, any.extract::(), @@ -1049,27 +1048,27 @@ fn tensor(py: Python<'_>, data: PyObject) -> PyResult { } #[pyfunction] -#[pyo3(signature = (shape, *, device=None), text_signature = "(shape:Sequence[int], device:Optional[Device]=None)")] +#[pyo3(signature = (*shape,device=None), text_signature = "(*shape:Shape, device:Optional[Device]=None)")] /// Creates a new tensor with random values. /// &RETURNS&: Tensor fn rand(_py: Python<'_>, shape: PyShape, device: Option) -> PyResult { let device = device.unwrap_or(PyDevice::Cpu).as_device()?; - let tensor = Tensor::rand(0f32, 1f32, shape.0, &device).map_err(wrap_err)?; + let tensor = Tensor::rand(0f32, 1f32, shape, &device).map_err(wrap_err)?; Ok(PyTensor(tensor)) } #[pyfunction] -#[pyo3(signature = (shape, *, device=None), text_signature = "(shape:Sequence[int], device:Optional[Device]=None)")] +#[pyo3(signature = (*shape,device=None), text_signature = "(*shape:Shape, device:Optional[Device]=None)")] /// Creates a new tensor with random values from a normal distribution. /// &RETURNS&: Tensor fn randn(_py: Python<'_>, shape: PyShape, device: Option) -> PyResult { let device = device.unwrap_or(PyDevice::Cpu).as_device()?; - let tensor = Tensor::randn(0f32, 1f32, shape.0, &device).map_err(wrap_err)?; + let tensor = Tensor::randn(0f32, 1f32, shape, &device).map_err(wrap_err)?; Ok(PyTensor(tensor)) } #[pyfunction] -#[pyo3(signature = (shape, *, dtype=None, device=None),text_signature = "(shape:Sequence[int], dtype:Optional[DType]=None, device:Optional[Device]=None)")] +#[pyo3(signature = (*shape, dtype=None, device=None),text_signature = "(*shape:Shape, dtype:Optional[DType]=None, device:Optional[Device]=None)")] /// Creates a new tensor filled with ones. /// &RETURNS&: Tensor fn ones( @@ -1083,12 +1082,12 @@ fn ones( Some(dtype) => PyDType::from_pyobject(dtype, py)?.0, }; let device = device.unwrap_or(PyDevice::Cpu).as_device()?; - let tensor = Tensor::ones(shape.0, dtype, &device).map_err(wrap_err)?; + let tensor = Tensor::ones(shape, dtype, &device).map_err(wrap_err)?; Ok(PyTensor(tensor)) } #[pyfunction] -#[pyo3(signature = (shape, *, dtype=None, device=None), text_signature = "(shape:Sequence[int], dtype:Optional[DType]=None, device:Optional[Device]=None)")] +#[pyo3(signature = (*shape, dtype=None, device=None), text_signature = "(*shape:Shape, dtype:Optional[DType]=None, device:Optional[Device]=None)")] /// Creates a new tensor filled with zeros. /// &RETURNS&: Tensor fn zeros( @@ -1102,7 +1101,7 @@ fn zeros( Some(dtype) => PyDType::from_pyobject(dtype, py)?.0, }; let device = device.unwrap_or(PyDevice::Cpu).as_device()?; - let tensor = Tensor::zeros(shape.0, dtype, &device).map_err(wrap_err)?; + let tensor = Tensor::zeros(shape, dtype, &device).map_err(wrap_err)?; Ok(PyTensor(tensor)) } diff --git a/candle-pyo3/src/shape.rs b/candle-pyo3/src/shape.rs new file mode 100644 index 0000000000..2668b7331b --- /dev/null +++ b/candle-pyo3/src/shape.rs @@ -0,0 +1,99 @@ +use ::candle::Tensor; +use pyo3::prelude::*; + +#[derive(Clone, Debug)] +/// Represents an absolute shape e.g. (1, 2, 3) +pub struct PyShape(Vec); + +impl<'source> pyo3::FromPyObject<'source> for PyShape { + fn extract(ob: &'source PyAny) -> PyResult { + if ob.is_none() { + return Err(PyErr::new::( + "Shape cannot be None", + )); + } + + let tuple = ob.downcast::()?; + if tuple.len() == 1 { + let first_element = tuple.get_item(0)?; + let dims: Vec = pyo3::FromPyObject::extract(first_element)?; + Ok(PyShape(dims)) + } else { + let dims: Vec = pyo3::FromPyObject::extract(tuple)?; + Ok(PyShape(dims)) + } + } +} + +impl From for ::candle::Shape { + fn from(val: PyShape) -> Self { + val.0.into() + } +} + +#[derive(Clone, Debug)] +/// Represents a shape with a hole in it e.g. (1, -1, 3) +pub struct PyShapeWithHole(Vec); + +impl<'source> pyo3::FromPyObject<'source> for PyShapeWithHole { + fn extract(ob: &'source PyAny) -> PyResult { + if ob.is_none() { + return Err(PyErr::new::( + "Shape cannot be None", + )); + } + + let tuple = ob.downcast::()?; + let dims: Vec = if tuple.len() == 1 { + let first_element = tuple.get_item(0)?; + pyo3::FromPyObject::extract(first_element)? + } else { + pyo3::FromPyObject::extract(tuple)? + }; + + // Ensure we have only positive numbers and at most one "hole" (-1) + let negative_ones = dims.iter().filter(|&&x| x == -1).count(); + let any_invalid_dimensions = dims.iter().any(|&x| x < -1 || x == 0); + if negative_ones > 1 || any_invalid_dimensions { + return Err(PyErr::new::(format!( + "Invalid dimension in shape: {:?}", + dims + ))); + } + + Ok(PyShapeWithHole(dims)) + } +} + +impl PyShapeWithHole { + /// Returns `true` if the shape is absolute e.g. (1, 2, 3) + pub fn is_absolute(&self) -> bool { + self.0.iter().all(|x| *x > 0) + } + + /// Convert a relative shape to an absolute shape e.g. (1, -1) -> (1, 12) + pub fn to_absolute(&self, t: &Tensor) -> PyResult { + if self.is_absolute() { + return Ok(PyShape( + self.0.iter().map(|x| *x as usize).collect::>(), + )); + } + + let mut elements = t.elem_count(); + let mut new_dims: Vec = vec![]; + for dim in self.0.iter() { + if *dim > 0 { + new_dims.push(*dim as usize); + elements /= *dim as usize; + } else if *dim == -1 { + new_dims.push(elements); + } else { + return Err(PyErr::new::(format!( + "Invalid dimension in shape: {}", + dim + ))); + } + } + Ok(PyShape(new_dims)) + } +} diff --git a/candle-pyo3/stub.py b/candle-pyo3/stub.py index 8e4318bcf5..336f674ba8 100644 --- a/candle-pyo3/stub.py +++ b/candle-pyo3/stub.py @@ -13,7 +13,7 @@ TYPING = """from typing import Any, Callable, Dict, List, Optional, Tuple, Union, Sequence from os import PathLike """ -CANDLE_SPECIFIC_TYPING = "from candle.typing import _ArrayLike, Device, Scalar, Index\n" +CANDLE_SPECIFIC_TYPING = "from candle.typing import _ArrayLike, Device, Scalar, Index, Shape\n" CANDLE_TENSOR_IMPORTS = "from candle import Tensor,DType,QTensor\n" RETURN_TYPE_MARKER = "&RETURNS&: " ADDITIONAL_TYPEHINTS = {} diff --git a/candle-pyo3/tests/native/test_shape.py b/candle-pyo3/tests/native/test_shape.py new file mode 100644 index 0000000000..864e24d679 --- /dev/null +++ b/candle-pyo3/tests/native/test_shape.py @@ -0,0 +1,31 @@ +from candle import Tensor +from candle import rand +import pytest + + +def test_absolute_shapes_are_valid(): + a = rand((10, 20)) + assert a.shape == (10, 20) + + b = rand(10, 20) + assert b.shape == (10, 20) + pytest.raises(OverflowError, lambda: rand((10, 20, -1))) + pytest.raises(OverflowError, lambda: rand(-1, 20)) + pytest.raises(TypeError, lambda: rand("foo", True)) + + +def test_relative_shapes_are_valid(): + a = rand(10, 20) + a = a.reshape((1, -1)) + assert a.shape == (1, 200) + + b = rand(10, 20) + b = b.reshape(-1, 1) + assert b.shape == (200, 1) + + c = rand(10, 20) + pytest.raises(TypeError, lambda: c.reshape(1, "foo")) + pytest.raises(ValueError, lambda: c.reshape(1, -2)) + pytest.raises(ValueError, lambda: c.reshape((-2, 1))) + pytest.raises(ValueError, lambda: c.reshape((0, 1))) + pytest.raises(ValueError, lambda: c.reshape((1, -1, -1)))