mlx_array_operations_with_jit_compilation_quickstart.py

1import mlx.core as mx
2
3# Create two arrays
4a = mx.array([1, 2, 3, 4])
5b = mx.array([5, 6, 7, 8])
6
7# Add them
8c = a + b
9
10# MLX is lazy, so c is not yet evaluated.
11# Calling print or mx.eval() will force evaluation.
12print(c)
13
14# Create a function and compile it
15def fast_fn(a, b):
16    return (a * b).sum()
17
18# Just-in-time compile the function
19g = mx.compile(fast_fn)
20
21# The first call will compile and execute
22# Subsequent calls will be faster
23result = g(a, b)
24print(result)