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

removed IntOrIntVec

Browse Source
This commit is contained in:
dc-dc-dc 2024-01-02 18:14:29 -05:00
parent 34108e780c
commit bd1a11bef7
2 changed files with 2 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
python/src/ops.cpp
View File

@ -3198,7 +3198,8 @@ void init_ops(py::module_& m) {
"tensordot", "tensordot",
[](const array& a, [](const array& a,
const array& b, const array& b,
const IntOrIntVec& dims, const std::variant<std::monostate, int, std::vector<std::vector<int>>>&
dims,
StreamOrDevice s) { StreamOrDevice s) {
if (std::holds_alternative<std::monostate>(dims)) { if (std::holds_alternative<std::monostate>(dims)) {
return tensordot(a, b, 2, s); return tensordot(a, b, 2, s);

2
python/src/utils.h
View File

@ -15,8 +15,6 @@ namespace py = pybind11;
using namespace mlx::core; using namespace mlx::core;
using IntOrVec = std::variant<std::monostate, int, std::vector<int>>; using IntOrVec = std::variant<std::monostate, int, std::vector<int>>;
using IntOrIntVec =
std::variant<std::monostate, int, std::vector<std::vector<int>>>;
using ScalarOrArray = std:: using ScalarOrArray = std::
variant<py::bool_, py::int_, py::float_, std::complex<float>, py::object>; variant<py::bool_, py::int_, py::float_, std::complex<float>, py::object>;
static constexpr std::monostate none{}; static constexpr std::monostate none{};