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

Fix CI format + build issue (#137)

Browse Source 
* fix ci

* Fix python bindings build

---------

Co-authored-by: Angelos Katharopoulos <a_katharopoulos@apple.com>
This commit is contained in:
Awni Hannun 2023-12-11 15:01:41 -08:00 committed by GitHub
parent 3214629601
commit b9226c367c
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 35 additions and 34 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

20
python/src/ops.cpp
View File

@ -1361,21 +1361,21 @@ void init_ops(py::module_& m) {
m.def( m.def(
"eye", "eye",
[](int n, [](int n,
py::object m_obj, std::optional<int> m,
py::object k_obj, int k,
Dtype dtype, std::optional<Dtype> dtype,
StreamOrDevice s) { StreamOrDevice s) {
int m = m_obj.is_none() ? n : m_obj.cast<int>(); return eye(n, m.value_or(n), k, dtype.value_or(float32), s);
int k = k_obj.is_none() ? 0 : k_obj.cast<int>();
return eye(n, m, k, dtype, s);
}, },
"n"_a, "n"_a,
"m"_a = py::none(), "m"_a = py::none(),
"k"_a = py::none(), "k"_a = 0,
"dtype"_a = std::nullopt, "dtype"_a = std::nullopt,
py::kw_only(), py::kw_only(),
"stream"_a = none, "stream"_a = none,
R"pbdoc( R"pbdoc(
eye(n: int, m: Optional[int] = None, k: int = 0, dtype: Optional[Dtype] = None, *, stream: Union[None, Stream, Device] = None) -> array
Create an identity matrix or a general diagonal matrix. Create an identity matrix or a general diagonal matrix.
Args: Args:
@ -1390,12 +1390,16 @@ void init_ops(py::module_& m) {
)pbdoc"); )pbdoc");
m.def( m.def(
"identity", "identity",
&identity, [](int n, std::optional<Dtype> dtype, StreamOrDevice s) {
return identity(n, dtype.value_or(float32), s);
},
"n"_a, "n"_a,
"dtype"_a = std::nullopt, "dtype"_a = std::nullopt,
py::kw_only(), py::kw_only(),
"stream"_a = none, "stream"_a = none,
R"pbdoc( R"pbdoc(
identity(n: int, dtype: Optional[Dtype] = None, *, stream: Union[None, Stream, Device] = None) -> array
Create a square identity matrix. Create a square identity matrix.
Args: Args:

3
python/tests/test_ops.py
View File

@ -1310,7 +1310,6 @@ class TestOps(mlx_tests.MLXTestCase):
b = mx.ones([2147484], mx.int8) b = mx.ones([2147484], mx.int8)
self.assertEqual((a + b)[0, 0].item(), 2) self.assertEqual((a + b)[0, 0].item(), 2)
def test_eye(self): def test_eye(self):
eye_matrix = mx.eye(3) eye_matrix = mx.eye(3)
np_eye_matrix = np.eye(3) np_eye_matrix = np.eye(3)
@ -1332,7 +1331,5 @@ class TestOps(mlx_tests.MLXTestCase):
self.assertTrue(np.array_equal(eye_matrix, np_eye_matrix)) self.assertTrue(np.array_equal(eye_matrix, np_eye_matrix))
if __name__ == "__main__": if __name__ == "__main__":
unittest.main() unittest.main()