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

remove monostate

Browse Source
This commit is contained in:
dc-dc-dc 2024-01-02 18:50:34 -05:00
parent bd1a11bef7
commit 3dc638bffb
1 changed files with 2 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
python/src/ops.cpp
View File

@ -3198,12 +3198,9 @@ void init_ops(py::module_& m) {
"tensordot", "tensordot",
[](const array& a, [](const array& a,
const array& b, const array& b,
const std::variant<std::monostate, int, std::vector<std::vector<int>>>& const std::variant<int, std::vector<std::vector<int>>>& dims,
dims,
StreamOrDevice s) { StreamOrDevice s) {
if (std::holds_alternative<std::monostate>(dims)) { if (auto pv = std::get_if<int>(&dims); pv) {
return tensordot(a, b, 2, s);
} else if (auto pv = std::get_if<int>(&dims); pv) {
return tensordot(a, b, *pv, s); return tensordot(a, b, *pv, s);
} else { } else {
return tensordot( return tensordot(