Wrong output when inf is the step in 'mlx.core.arange'

[Redmept1on]

Describe the bug According to arange's calculation method (Generate numbers in the half open interval [start, stop)), when step is inf, the generated array should contain the start value, but mlx.core.arange did not achieve this.

To Reproduce

Include code snippet

import mlx.core as mx
import numpy as np

start = 0.1
stop = 2
step = float('inf')

#mlx
range = mx.arange(start, stop, step)
print('mlx:', range)

#mlx
np_range = np.arange(start, stop, step)
print('numpy:', np_range)

Expected behavior array([0.1], dtype=float32)

Desktop (please complete the following information):

• OS Version: MacOS 14.2.1
• Version 0.10.0

[nakranivaibhav]

@awni I can take this up

[awni]

Great!

[nakranivaibhav]

@awni Can you help me how to setup an editable version of mlx on my local machine. I am getting errors(Tried many times). I want to test the code as I make changes along side. Or is there a better way to test it?

Steps done so far: -> Created a venv -> pip installed pybind11 separately. (For som reason its not picking up the module from pyproject.toml) -> Added the installation prefix of “pybind11” to CMAKE_PREFIX_PATH.

[awni]

Check out the docs: https://ml-explore.github.io/mlx/build/html/install.html#python-api

You need to install pybind11 the way it is mentioned there.

[nakranivaibhav]

alright thanks for pointing it out. I will try that

[nakranivaibhav]

@awni What might be a good way to test the code changes. I have built using the -e flag but the c code that I have changed does not reflect in the library where I have installed for testing. Do I have to rebuild the entire library everytime I change the c or cpp code. Or is there a better way.

P.S: Sorry to bug you so much. This is the first time I am working with binded python code.

[awni]

Yes, you have to rebuild the C++ code if you change it :). Cmake usually does a good job of only building parts which are touched so it shouldn't take that long.

[noahfarr]

I can pick this up

[awni]

Thanks!! Assigned to you!

[noahfarr]

Just to be clear on how the behavior should be: If step is inf it should return an array with the start value. But what about when start or stop are inf? What about when step is -inf? I feel like it would be most intuitive to throw an error in any case since they're all basically undefined behaviors.

[awni]

Could you check what NumPy does for those cases? Typically we just follow what they do for these random edge cases unless there is a reason not to (e.g. it makes it more complex or it's a lot more intuitive to do something different).

[noahfarr]

I checked but they do not mention anything about this. I also checked the torch docs and there was no information there either. https://numpy.org/doc/stable/reference/generated/numpy.arange.html https://pytorch.org/docs/stable/generated/torch.arange.html

[awni]

You usually have to use code to check as the docs won't tell you every edge case:

>>> np.arange(0, float("inf"), 1)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
ValueError: Maximum allowed size exceeded

[noahfarr]

Numpy throws an error for start/stop inf: b = np.arange(0, float("inf"), 1) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ ValueError: Maximum allowed size exceeded

but for step inf it returns an array with the start value

[awni]

That is probably what we should do as well. In fact we might want to put a size limit of INT_MAX since that's the largest 1D array we can make right now

[noahfarr]

Do you think the solution from my PR is fine?

[nakranivaibhav]

Do let me know how did you test the code. Did you rebuild everything. @awni sorry for not working on this. I even have the diff code but I couldn't test it even after rebuilding.

[noahfarr]

I simply followed the instructions here: Documentation