Added linspace (#181)

* linspace ops support

---------

Co-authored-by: Awni Hannun <awni@apple.com>
This commit is contained in:
Abe Leininger
2023-12-18 19:57:55 -08:00
committed by GitHub
parent f4f6e17d45
commit e6872a4149
6 changed files with 92 additions and 1 deletions

View File

@@ -1184,6 +1184,32 @@ void init_ops(py::module_& m) {
This can lead to unexpected results for example if `start + step`
is a fractional value and the `dtype` is integral.
)pbdoc");
m.def(
"linspace",
[](Scalar start, Scalar stop, int num, Dtype dtype, StreamOrDevice s) {
return linspace(
scalar_to_double(start), scalar_to_double(stop), num, dtype, s);
},
"start"_a,
"stop"_a,
"num"_a = 50,
"dtype"_a = float32,
"stream"_a = none,
R"pbdoc(
linspace(start, stop, num: Optional[int] = 50, dtype: Optional[Dtype] = float32, stream: Union[None, Stream, Device] = None) -> array
Generate ``num`` evenly spaced numbers over interval ``[start, stop]``.
Args:
start (scalar): Starting value.
stop (scalar): Stopping value.
num (int, optional): Number of samples, defaults to ``50``.
dtype (Dtype, optional): Specifies the data type of the output,
default to ``float32``.
Returns:
array: The range of values.
)pbdoc");
m.def(
"take",
[](const array& a,