mirror of
https://github.com/ml-explore/mlx.git
synced 2025-10-18 15:28:16 +08:00
Minor updates to address a few issues (#537)
* docs on arg indices return type * arange with nan * undo isort
This commit is contained in:
@@ -2254,7 +2254,7 @@ void init_ops(py::module_& m) {
|
||||
singleton dimensions, defaults to `False`.
|
||||
|
||||
Returns:
|
||||
array: The output array with the indices of the minimum values.
|
||||
array: The ``uint32`` array with the indices of the minimum values.
|
||||
)pbdoc");
|
||||
m.def(
|
||||
"argmax",
|
||||
@@ -2287,7 +2287,7 @@ void init_ops(py::module_& m) {
|
||||
singleton dimensions, defaults to `False`.
|
||||
|
||||
Returns:
|
||||
array: The output array with the indices of the maximum values.
|
||||
array: The ``uint32`` array with the indices of the maximum values.
|
||||
)pbdoc");
|
||||
m.def(
|
||||
"sort",
|
||||
@@ -2343,7 +2343,7 @@ void init_ops(py::module_& m) {
|
||||
If unspecified, it defaults to -1 (sorting over the last axis).
|
||||
|
||||
Returns:
|
||||
array: The indices that sort the input array.
|
||||
array: The ``uint32`` array containing indices that sort the input.
|
||||
)pbdoc");
|
||||
m.def(
|
||||
"partition",
|
||||
@@ -2416,7 +2416,7 @@ void init_ops(py::module_& m) {
|
||||
If unspecified, it defaults to ``-1``.
|
||||
|
||||
Returns:
|
||||
array: The indices that partition the input array.
|
||||
array: The `uint32`` array containing indices that partition the input.
|
||||
)pbdoc");
|
||||
m.def(
|
||||
"topk",
|
||||
|
@@ -980,6 +980,17 @@ class TestOps(mlx_tests.MLXTestCase):
|
||||
self.assertEqual(z.tolist(), [5, 6, 7])
|
||||
|
||||
def test_arange_overload_dispatch(self):
|
||||
with self.assertRaises(ValueError):
|
||||
a = mx.arange(float("nan"), 1, 5)
|
||||
with self.assertRaises(ValueError):
|
||||
a = mx.arange(0, float("nan"), 5)
|
||||
with self.assertRaises(ValueError):
|
||||
a = mx.arange(0, 2, float("nan"))
|
||||
with self.assertRaises(ValueError):
|
||||
a = mx.arange(0, float("inf"), float("inf"))
|
||||
with self.assertRaises(ValueError):
|
||||
a = mx.arange(float("inf"), 1, float("inf"))
|
||||
|
||||
a = mx.arange(5)
|
||||
expected = [0, 1, 2, 3, 4]
|
||||
self.assertListEqual(a.tolist(), expected)
|
||||
|
Reference in New Issue
Block a user