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

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)

image

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

Desktop (please complete the following information):

  • OS Version: MacOS 14.2.1
  • Version 0.10.0

About this issue

  • Original URL
  • State: closed
  • Created 5 months ago
  • Comments: 18 (13 by maintainers)

Most upvoted comments

Great!

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 I can take this up