mirror of
https://github.com/ml-explore/mlx.git
synced 2025-09-11 06:24:35 +08:00
Add inner / outer op (#348)
* inner / outer impl * python tests * ops list and ack * updated descriptions * use test helper * removed dtype check and flatten outer to 1-D * updated docs * just use the reshape to flatten
This commit is contained in:
@@ -3250,4 +3250,46 @@ void init_ops(py::module_& m) {
|
||||
Returns:
|
||||
result (array): The tensor dot product.
|
||||
)pbdoc");
|
||||
|
||||
m.def(
|
||||
"inner",
|
||||
&inner,
|
||||
"a"_a,
|
||||
"b"_a,
|
||||
py::pos_only(),
|
||||
py::kw_only(),
|
||||
"stream"_a = none,
|
||||
R"pbdoc(
|
||||
inner(a: array, b: array, /, *, stream: Union[None, Stream, Device] = None) -> array
|
||||
|
||||
Ordinary inner product of vectors for 1-D arrays, in higher dimensions a sum product over the last axes.
|
||||
|
||||
Args:
|
||||
a (array): Input array
|
||||
b (array): Input array
|
||||
|
||||
Returns:
|
||||
result (array): The inner product.
|
||||
)pbdoc");
|
||||
|
||||
m.def(
|
||||
"outer",
|
||||
&outer,
|
||||
"a"_a,
|
||||
"b"_a,
|
||||
py::pos_only(),
|
||||
py::kw_only(),
|
||||
"stream"_a = none,
|
||||
R"pbdoc(
|
||||
outer(a: array, b: array, /, *, stream: Union[None, Stream, Device] = None) -> array
|
||||
|
||||
Compute the outer product of two 1-D arrays, if the array's passed are not 1-D a flatten op will be run beforehand.
|
||||
|
||||
Args:
|
||||
a (array): Input array
|
||||
b (array): Input array
|
||||
|
||||
Returns:
|
||||
result (array): The outer product.
|
||||
)pbdoc");
|
||||
}
|
||||
|
Reference in New Issue
Block a user