[BUG] linspace should always include the stop value

[acsweet]

Describe the bug
linspace truncates values in some cases

To Reproduce

(lazily casting to numpy to print to see precision)

import mlx.core as mx
import jax.numpy as jnp
import numpy as np

range = (0.21436651, 0.96568555)
num = 6
a = mx.linspace(range[0], range[1], num, dtype=mx.float32)
print(np.array(a))
a = jnp.linspace(range[0], range[1], num, dtype=jnp.float32)
print(np.array(a))
# [0.21436651 0.3646303  0.5148941  0.6651579  0.8154217  0.9656855 ]
# [0.21436651 0.3646303  0.5148941  0.6651579  0.81542176 0.96568555]

Expected behavior
I'd expect linspace to produce the same results as jax, without truncation for the same dtype

Desktop (please complete the following information):

• OS Version: 15.2
• Version: 0.23.0

Additional context
I was using this to create bin edges for a histogram, and it was excluding the right most value because of the truncation when calculating counts. I can work around this by setting the last value if it's not an easy fix.

[angeloskath]

This is normal floating point errors. The computation on the MLX side is happening in float32. I think (but haven't verified) that in numpy even though float32 is the type of the result, internally a double is used. Just to make the point clearer, 0.96568555 is actually not even representable in float32.

Another way of looking at it is that if we write the computation in float32 both frameworks get the exact same result (as expected)

import mlx.core as mx
import numpy as np

range = (0.21436651, 0.96568555)
num = 6
a = mx.linspace(0, 1, num, dtype=mx.float32) * (range[1] - range[0]) + range[0]
b = np.linspace(0, 1, num, dtype=np.float32) * (range[1] - range[0]) + range[0]
print(b-a)
# [0. 0. 0. 0. 0. 0.]

print(b - np.linspace(range[0], range[1], num, dtype=np.float32))
# [0.0000000e+00 2.9802322e-08 0.0000000e+00 5.9604645e-08 0.0000000e+00 0.0000000e+00]

[acsweet]

@angeloskath Thank you!
I think my question was malformed, and I apologize.

Taking the closest fp32 representation near 0.96568555, I think is 0.9656855463981628, and computing linspace with that yields a different final value for me. Apologies too, I was comparing jax and mlx.
Note that it captures the end value when num=5, is it just an implementation detail of linspace for mlx and how it partitions the spaces? I think I expected it to be inclusive of both edges (similar to jax).

import mlx.core as mx
import jax.numpy as jnp

num = 6
range_bin = (0.21436651051044464, 0.9656855463981628)
a = mx.linspace(range_bin[0], range_bin[1], num, dtype=mx.float32)
print(a[0] == range_bin[0], a[-1] == range_bin[-1])
print(a[0].item(), a[-1].item())
# array(True, dtype=bool) array(False, dtype=bool)
# 0.21436651051044464 0.9656854867935181

b = jnp.linspace(range_bin[0], range_bin[1], num, dtype=jnp.float32)
print(b[0] == range_bin[0], b[-1] == range_bin[-1])
print(b[0].item(), b[-1].item())
# True True
# 0.21436651051044464 0.9656855463981628

[acsweet]

I went and checked what numpy and jax are doing!

For numpy, it seems like it just explicitly sets the last value(s) to address cases likes this
https://github.com/numpy/numpy/blob/v2.1.0/numpy/_core/function_base.py#L180-L183

Jax actually (kind of "cheats") generates all points excluding the end, then appends the last point(s) (if endpoint=True)
https://github.com/jax-ml/jax/blob/main/jax/_src/numpy/lax_numpy.py#L6539-L6543

I think it's entirely up to you, if you'd like to also have guaranteed inclusive behavior of both ends (or maybe a docs note if not so?).

[angeloskath]

No need for apologies. I am reopening it and I will look tomorrow into how easy it would be to ensure inclusivity.

[acsweet]

Thank you!