MLX is faster than torch and numpy for converting arrays to numpy. Why?

[famishedrover]

I am benchmarking conversion of framework specific array type to numpy array and generally found that mlx is 2-3x better than torch to numpy conversion (which is amazing), however, what surprised me was that a numpy.array(x) call when x is torch.Tensor is similar to when x is ndarray. Therefore, mlx.core.array -> ndarray is 2-3x faster than ndarray -> ndarray!

MLX Array -> Numpy

v = [[], [], []]
for ix in range(10):
    start = time.time()
    a = mx.random.uniform(1, 1000, (10000, 10000))
    mx.eval(a)
    second = time.time()
    b = np.array(a)
    end = time.time()
    v[0].append(second - start)
    v[1].append(end - second)
    v[2].append(end - start)

v = np.array(v).mean(axis=-1)
print(f"Creation = {v[0]}, Conversion = {v[1]}, Total = {v[2]}")

Creation = 1.317774200439453, Conversion = 0.13875243663787842, Total = 1.4565266370773315

(and the following when I used mx.random.randint instead of float from uniform)

Creation = 1.4690469980239869, Conversion = 0.14126429557800294, Total = 1.6103112936019897

Torch.Tensor CPU -> Numpy

v = [[], [], []]
for ix in range(10):
    start = time.time()
    a = torch.randint(1, 1000, (10000, 10000))
    second = time.time()
    b = np.array(a)
    end = time.time()
    v[0].append(second - start)
    v[1].append(end - second)
    v[2].append(end - start)

v = np.array(v).mean(axis=-1)
print(f"Creation = {v[0]}, Conversion = {v[1]}, Total = {v[2]}")

Creation = 1.6754926204681397, Conversion = 0.2948038101196289, Total = 1.9702964305877686

Torch Tensor GPU -> Numpy

v = [[], [], []]
for ix in range(10):
    start = time.time()
    a = torch.randint(1, 1000, (10000, 10000)).to('cuda')
    second = time.time()
    b = np.array(a.cpu())
    end = time.time()
    v[0].append(second - start)
    v[1].append(end - second)
    v[2].append(end - start)

v = np.array(v).mean(axis=-1)
print(f"Creation = {v[0]}, Conversion = {v[1]}, Total = {v[2]}")

Creation = 2.0264783620834352, Conversion = 0.9157697916030884, Total = 2.9422481536865233

Numpy -> Numpy

v = [[], [], []]
for ix in range(10):
    start = time.time()
    a = np.random.randint(1, 1000, (10000, 10000))
    second = time.time()
    b = np.array(a)
    end = time.time()
    v[0].append(second - start)
    v[1].append(end - second)
    v[2].append(end - start)

v = np.array(v).mean(axis=-1)
print(f"Creation = {v[0]}, Conversion = {v[1]}, Total = {v[2]}")

Creation = 0.5674269914627075, Conversion = 0.2816884517669678, Total = 0.8491154432296752

Jax -> np array

v = [[], [], []]
key = jax.random.PRNGKey(0)

for ix in range(10):
    start = time.time()
    a = jax.random.randint(key, (10000, 10000), minval=1, maxval=1000)
    second = time.time()
    b = np.array(a)
    end = time.time()
    v[0].append(second - start)
    v[1].append(end - second)
    v[2].append(end - start)

v = np.array(v).mean(axis=-1)
print(f"Creation = {v[0]}, Conversion = {v[1]}, Total = {v[2]}")

Creation = 3.48213586807251, Conversion = 0.15759737491607667, Total = 3.6397332429885862

[awni]

I got this on my M1 Max with the latest MLX (using Nanobind is a lot faster), also using zeros instead of random:

Numpy: Creation = 0.007582283020019532, Conversion = 0.1514657258987427, Total = 0.1590480089187622
MLX: Creation = 0.006192994117736816, Conversion = 0.04600672721862793, Total = 0.052199721336364746

[awni]

I really couldn't say exactly why it's so much faster... when you do np.array(a) and a is an MLX array, then we just create a memoryview which gets used by NumPy to construct the actual array. So indeed it is quite strange that numpy -> numpy is slower than MLX -> numpy..

[famishedrover]

I speculate whether it has to do with how mlx array are allocated memory v/s how numpy allocates it.