In-Place operators (+=, -=) for MLXArray

[kemchenj]

I have been using mlx-swift for my projects and have found it to be incredibly useful. However, I noticed that MLXArray does not support in-place operators such as += and -=.

Is there a specific reason why these operators have not been implemented? Additionally, are there any plans to add these in-place operators in the future?

These operators would be very beneficial for performance-critical applications where minimizing memory allocations and copies is crucial.

Thank you for your time and for the great work on mlx-swift!

[davidkoski]

Although they can be added syntactically I don't know if there is really such a thing under the hood.

@awni are there any in-place operations in mlx core?

[awni]

Right from the user standpoint there are in-place operations. So if you do an in place operation then the object is unchanged:

orig_id = id(a)
a += b
new_id = id(a)

Then orig_id == new_id.

However, when it comes to memory management, that is completely opaque to the user. MLX often re-uses memory regardless of if the op is in place or not. It will reuse memory when it is safe to do so (typically if you don't retain references to arrays and the buffer sizes match).

So in the following:

def fun():
  a = mx.array([1,2,3])
  a = a + mx.array([2])
  return a

out = fun()
mx.eval(out)

Even though the op is not in-place since the reference to the original array is out of scope before the call to eval the memory buffer is likely reused

[kemchenj]

Thank you for the detailed explanation. I understand that these operators may not significantly impact memory management due to MLX's efficient handling of memory reuse. However, from a syntactical perspective, having in-place operators like += and -= would still be valuable.

Additionally, aligning these APIs with the Python version, which implements iadd/isub/..., could help developers transition from Python to Swift more smoothly, reducing potential glitches.

[davidkoski]

Agreed, these += and friends operators should be added.