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:
Diogo
2024-01-07 12:01:09 -05:00
committed by GitHub
parent 6ea6b4258d
commit 449b43762e
7 changed files with 140 additions and 5 deletions

View File

@@ -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");
}