zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-08-21 12:06:42 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

added pbdoc

Browse Source
This commit is contained in:
dc-dc-dc 2024-01-02 13:20:47 -05:00
parent 3fb3385a1a
commit eb696110d8
1 changed files with 17 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

18
python/src/ops.cpp
View File

@ -3215,5 +3215,21 @@ void init_ops(py::module_& m) {
"dims"_a = 2, "dims"_a = 2,
py::kw_only(), py::kw_only(),
"stream"_a = none, "stream"_a = none,
""); R"pbdoc(
tensordot(a: array, b: array, /, dims: Union[int, List[List[int]]] = 2, *, stream: Union[None, Stream, Device] = None) -> array
Compute the tensor dot product along the specified axes.
Args:
a (array): Input array
b (array): Input array
dims (int or list(list(int)), optional): The number of dimensions to
sum over. If an integer is provided, then sum over the last
``dims`` dimensions of ``a`` and the first ``dims`` dimensions of
``b``. If a list of lists is provided, then sum over the
corresponding dimensions of ``a`` and ``b``. (default: 2)
Returns:
result (array): The tensor dot product.
)pbdoc");
} }