Add support for repeat #278

Bahaatbb · 2023-12-24T08:48:14Z

Proposed changes

resolves issue #258

Add op:

mx.repeat

same behavior as np.repeat

Checklist

Put an x in the boxes that apply.

I have read the CONTRIBUTING document
I have run pre-commit run --all-files to format my code / installed pre-commit prior to committing changes
I have added tests that prove my fix is effective or that my feature works
I have updated the necessary documentation (if needed)

awni

Thanks for adding this! I left a few minor comments. But mainly, the core implementation needs to change a bit as the way you have it, the number of arrays/ops scales with the size of the array which won't work well. I left a comment with an example of what I mean. Let me know if it's not clear.

awni · 2023-12-27T05:51:50Z

mlx/ops.cpp

@@ -1,11 +1,10 @@
 // Copyright © 2023 Apple Inc.

+#include "mlx/ops.h"


Why move this?

sorry, this seems like a slip

awni · 2023-12-27T05:54:45Z

tests/ops_tests.cpp

@@ -2233,3 +2233,40 @@ TEST_CASE("test quantize dequantize") {
    CHECK(max_diff <= 127.0 / (1 << i));
  }
 }
+
+TEST_CASE("repeat test with axis") {


Just call this test repeat (since it also checks without axis)

awni · 2023-12-27T05:57:37Z

mlx/ops.cpp

+}
+
+array repeat(const array& arr, int repeats, StreamOrDevice s) {
+  return flatten(repeat(arr, repeats, arr.ndim() - 1, s));


Add the stream to the call to flatten. Also I would flatten first then repeat (with axis 0)..since you have to flatten anyway seems slightly more efficient to flatten first.

awni · 2023-12-27T06:08:54Z

mlx/ops.cpp

+  std::vector<array> arrays_to_concat;
+  arrays_to_concat.reserve(repeats * arr.shape(axis));
+
+  for (int i = 0; i < arr.shape(axis); ++i) {
+    std::vector<int> start_indices(arr.ndim(), 0);
+    std::vector<int> stop_indices = arr.shape();
+    start_indices[axis] = i;
+    stop_indices[axis] = i + 1;
+    for (int j = 0; j < repeats; ++j) {
+      arrays_to_concat.push_back(slice(arr, start_indices, stop_indices, s));
+    }
+  }


This needs to be implemented with a reshape and a concatenate that is order repeats and critically does not scale with the shape of the repeating axis. For example if you are repeating:

a = array([[0, 1], [2, 3]])

along axis 0 (in numpy like psuedo-code) do:

concatenate([a]*repeats, -1).reshape(-1, a.shape[1])

and along axis 1 do:

stack([a]*repeats, -1).reshape(a.shape[0], -1)

hey just to wrap my head around this, do you mean I should handle the cases of 0 and 1 axis repetitions alone and the other axis with the same logic or for the other axis with only a reshape and a concatenate that is order repeats

Yea, what I'm saying is we should avoid ever making repeats * array.shape(axis) subarrays to concatenate using something like the strategy I outlined above for a 2D array which just makes repeats arrays. I think it should generalize to ND arrays but let me know if I'm missing something.

Yes we definitely can generalize this for ND array using something like this for axis 0 or 1

def repeat(arr, repeats, axis): new_shape = np.array(arr.shape) new_shape[axis] *= repeats if axis == 0: repeated = np.concatenate([arr]*repeats, axis=-(len(arr.shape)-1)).flatten() return repeated.reshape(new_shape) elif axis == 1: repeated = np.stack([arr for _ in range(repeats)], axis=-(len(arr.shape)-1)).flatten() return repeated.reshape(new_shape)

but what I'm trying to understand is if you think there is a better way for other than making slices for axis bigger than 0 or 1 for bigger arrays.

Yea exactly, I think it should work for any axis of any n-darray. You have to concatenate along the correct dimension followed by a reshape.

Hey please check the new commit it should be exactly what you suggested

Much better thanks!! I left a few more comments. I think we should be good to go after you address them.

awni · 2023-12-27T18:09:58Z

mlx/ops.cpp

+    repeated_arrays.push_back(expand_dims(arr, -1, s));
+  }
+  array repeated =
+      flatten(concatenate(repeated_arrays, -1 * concat_axes, s), s);


Why do you flatten the array here?

Also I think it's cleaner if you replace -1 * concate_axes with axis + 1.

Why do you flatten the array here?

Oh, thank you for noticing we don't really need it here,
I first implement it in python and didn't realize that I added the dim by hand, but here we have expand_dims.
Fixed it, what a keen eye 🥇

awni · 2023-12-27T18:10:34Z

python/src/ops.cpp

+      R"pbdoc(
+      repeat(array: array, repeats: int, axis: Optional[int] = None, *, stream: Union[None, Stream, Device] = None) -> array
+
+      Repeate an array along a specified axis


nit period after this.

awni · 2023-12-27T18:10:53Z

python/src/ops.cpp

+      Args:
+          array (array): Input array.
+          repeats (int): The number of repetitions for each element.
+          axis (int, optional): The axis in which to repeat the array along. Defaults to ``None``.


Can you say more about the default behavior here?

awni · 2023-12-27T18:14:45Z

mlx/ops.cpp

+  if (repeats <= 0) {
+    std::vector<int> new_shape(arr.shape());
+    new_shape[axis] = repeats > 0 ? repeats : 0;
+    return zeros(new_shape, arr.dtype(), s);
+  }


Throw on negative repeats (as in NumPy) and add a test case to check that it throws.

For repeats==0 return an empty array with just one dimensions. Which I think you can do with array({}, arr.dtype())

awni · 2023-12-27T18:15:17Z

_deps/doctest-src

Could you remove this. I think you added by mistake.

am not sure what that is, but always happy to remove things 😅

awni · 2023-12-27T19:39:41Z

mlx/ops.cpp

+  axis = normalize_axis(axis, arr.ndim());
+
+  if (repeats < 0) {
+    throw std::invalid_argument("Number of repeats cannot be negative");


One more nit: add "[repeat]" to the error message

awni · 2023-12-27T19:40:37Z

python/src/ops.cpp

+          repeats (int): The number of repetitions for each element.
+          axis (int, optional): The axis in which to repeat the array along. If
+            unspecified it uses the flattened array of the input and repeates 
+            along the 0 axis.


nit: "along axis 0"

awni

Thanks for adding this! Looks great. Could you fix the last two nits and then I will merge it?

* add repeat function * fix styling * optimizing repeat * fixed minor issues * not sure why that folder is there xD * fixed now for sure * test repeat not repeat test * Fixed --------- Co-authored-by: Bahaa Eddin tabbakha <bahaa@Bahaas-MacBook-Pro.local>

Bahaa Eddin tabbakha added 2 commits December 24, 2023 11:37

add repeat function

c9328d5

fix styling

93ea457

dc-dc-dc mentioned this pull request Dec 26, 2023

Support for einops #172

Closed

awni requested changes Dec 27, 2023

View reviewed changes

optimizing repeat

91488e2

awni reviewed Dec 27, 2023

View reviewed changes

Bahaa Eddin tabbakha added 4 commits December 27, 2023 21:44

fixed minor issues

dc19971

not sure why that folder is there xD

2118df8

fixed now for sure

bfd1355

test repeat not repeat test

04fccee

Bahaatbb requested a review from awni December 27, 2023 19:36

awni reviewed Dec 27, 2023

View reviewed changes

awni approved these changes Dec 27, 2023

View reviewed changes

Fixed

6867e6f

awni merged commit ff2b58e into ml-explore:main Dec 27, 2023

awni mentioned this pull request Jan 2, 2024

Alternative to np.repeat? #258

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add support for repeat #278

Add support for repeat #278

Bahaatbb commented Dec 24, 2023

awni left a comment

awni Dec 27, 2023

Bahaatbb Dec 27, 2023

awni Dec 27, 2023

awni Dec 27, 2023

awni Dec 27, 2023

Bahaatbb Dec 27, 2023 •

edited

Loading

awni Dec 27, 2023

Bahaatbb Dec 27, 2023

awni Dec 27, 2023

Bahaatbb Dec 27, 2023

awni Dec 27, 2023

awni Dec 27, 2023

Bahaatbb Dec 27, 2023

awni Dec 27, 2023

awni Dec 27, 2023

awni Dec 27, 2023

awni Dec 27, 2023

Bahaatbb Dec 27, 2023

awni Dec 27, 2023

Bahaatbb Dec 27, 2023

awni Dec 27, 2023

awni left a comment

		@@ -1,11 +1,10 @@
		// Copyright © 2023 Apple Inc.

		#include "mlx/ops.h"

Add support for repeat #278

Add support for repeat #278

Conversation

Bahaatbb commented Dec 24, 2023

Proposed changes

Checklist

awni left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Bahaatbb Dec 27, 2023 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

awni left a comment

Choose a reason for hiding this comment

Bahaatbb Dec 27, 2023 •

edited

Loading