zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-10-18 15:28:16 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

implemented Flatten Module (#149)

Browse Source 
* implemented flatten op

---------

Co-authored-by: Awni Hannun <awni@apple.com>
This commit is contained in:
__mo_san__
2023-12-17 06:54:37 +01:00
committed by GitHub
parent eebd7c275d
commit 52e1589a52
8 changed files with 113 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

15
python/src/array.cpp
View File

@@ -728,6 +728,21 @@ void init_array(py::module_& m) {
return power(a, to_array(v, a.dtype()));
},
"other"_a)
.def(
"flatten",
[](const array& a,
int start_axis,
int end_axis,
const StreamOrDevice& s) {
return flatten(a, start_axis, end_axis);
},
"start_axis"_a = 0,
"end_axis"_a = -1,
py::kw_only(),
"stream"_a = none,
R"pbdoc(
See :func:`flatten`.
)pbdoc")
.def(
"reshape",
[](const array& a, py::args shape, StreamOrDevice s) {

27
python/src/ops.cpp
View File

@@ -61,6 +61,33 @@ void init_ops(py::module_& m) {
Returns:
array: The reshaped array.
)pbdoc");
m.def(
"flatten",
[](const array& a,
int start_axis,
int end_axis,
const StreamOrDevice& s) { return flatten(a, start_axis, end_axis); },
"a"_a,
py::pos_only(),
"start_axis"_a = 0,
"end_axis"_a = -1,
py::kw_only(),
"stream"_a = none,
R"pbdoc(
flatten(a: array, /, start_axis: int = 0, end_axis: int = -1, *, stream: Union[None, Stream, Device] = None) -> array
Flatten an array.
Args:
a (array): Input array.
start_axis (int, optional): The first dimension to flatten. Defaults to ``0``.
end_axis (int, optional): The last dimension to flatten. Defaults to ``-1``.
stream (Stream, optional): Stream or device. Defaults to ``None``
in which case the default stream of the default device is used.
Returns:
array: The flattened array.
)pbdoc");
m.def(
"squeeze",
[](const array& a, const IntOrVec& v, const StreamOrDevice& s) {

9
python/tests/test_ops.py
View File

@@ -1426,6 +1426,15 @@ class TestOps(mlx_tests.MLXTestCase):
np_c = np.stack([np_a, np_b], axis=1)
self.assertTrue(np.array_equal(c, np_c))
def test_flatten(self):
x = mx.zeros([2, 3, 4])
self.assertEqual(mx.flatten(x).shape, [2 * 3 * 4])
self.assertEqual(mx.flatten(x, start_axis=1).shape, [2, 3 * 4])
self.assertEqual(mx.flatten(x, end_axis=1).shape, [2 * 3, 4])
self.assertEqual(x.flatten().shape, [2 * 3 * 4])
self.assertEqual(x.flatten(start_axis=1).shape, [2, 3 * 4])
self.assertEqual(x.flatten(end_axis=1).shape, [2 * 3, 4])
if __name__ == "__main__":
unittest.main()