zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-11-04 10:38:10 +08:00
Code Issues Actions 1 Packages Projects Releases Wiki Activity

Added support for atleast_1d, atleast_2d, atleast_3d (#694)

Browse Source
This commit is contained in:
Hinrik Snær Guðmundsson
2024-02-19 12:40:52 -05:00
committed by GitHub
parent e1bdf6a8d9
commit f883fcede0
7 changed files with 241 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

60
python/src/ops.cpp
View File

@@ -3636,4 +3636,64 @@ void init_ops(py::module_& m) {
Returns:
array: The extracted diagonal or the constructed diagonal matrix.
)pbdoc");
m.def(
"atleast_1d",
&atleast_1d,
"a"_a,
py::pos_only(),
py::kw_only(),
"stream"_a = none,
R"pbdoc(
atleast_1d(a: array, stream: Union[None, Stream, Device] = None) -> array
Convert array to have at least one dimension.
args:
a (array): Input array
stream (Union[None, Stream, Device], optional): The stream to execute the operation on.
Returns:
array: An array with at least one dimension.
)pbdoc");
m.def(
"atleast_2d",
&atleast_2d,
"a"_a,
py::pos_only(),
py::kw_only(),
"stream"_a = none,
R"pbdoc(
atleast_2d(a: array, stream: Union[None, Stream, Device] = None) -> array
Convert array to have at least two dimensions.
args:
a (array): Input array
stream (Union[None, Stream, Device], optional): The stream to execute the operation on.
Returns:
array: An array with at least two dimensions.
)pbdoc");
m.def(
"atleast_3d",
&atleast_3d,
"a"_a,
py::pos_only(),
py::kw_only(),
"stream"_a = none,
R"pbdoc(
atleast_3d(a: array, stream: Union[None, Stream, Device] = None) -> array
Convert array to have at least three dimensions.
args:
a (array): Input array
stream (Union[None, Stream, Device], optional): The stream to execute the operation on.
Returns:
array: An array with at least three dimensions.
)pbdoc");
}