April 2, 2024

What’s new in Mojo 24.2: Mojo Nightly, Enhanced Python Interop, OSS stdlib and more

Shashank Prasanna

Last week we made MAX 24.2 available for download, and we shared more details in the MAX 24.2 announcement and the Mojo open-source blog posts. In this blog post. I’ll dive a little deeper into all the new features specifically in Mojo🔥. This will be your example-driven guide to Mojo SDK 24.2, as part of the latest MAX release. If I had to pick a name for this release, I’d call it MAXimum⚡ Mojo🔥 Momentum 🚀 because there is so much much good stuff in this release, particularly for Python developers, adopting Mojo. 

Here’s a summary of all the key features in this release.

Through the rest of the blog post, I’ll share each feature listed above with code examples that you can copy/paste and follow along. You can also access all the code samples in this blog post in a Jupyter Notebook on GitHub. As always, the official changelog has an exhaustive list of new features, what’s changed, what’s removed, and what’s fixed. And before we continue Don’t forget to upgrade your Mojo🔥. If you haven’t already, head over to the getting started page and follow the 3 simple steps to update Mojo to the latest release. Let’s dive into the new features.

New Mojo standard library features

Starting with this release, Mojo standard library is now open-source! You can head over to Mojo repository on GitHub where you’ll find the source code and contributing guide. Here are the new features in the Mojo standard library, explained with code examples.

DynamicVector is now called List

I had previously ignored using DynamicVector because I didn’t realize it’s similar to Python’s list data container. Starting with this release, DynamicVector has been renamed to List. If you are a Python list user, then you should feel right at home with Mojo List. Let’s take a look at a simple stack-based sorting algorithm using Lists in Mojo.

Mojo
from random import random_si64 def stack_sort(list: List[Int]) -> List[Int]: stack = List[Int]() input_list = List[Int](list) sorted_list = List[Int]() while len(input_list): temp = input_list.pop_back() while len(stack) and stack[-1] < temp: input_list.append(stack.pop_back()) stack.append(temp) while len(stack): sorted_list.append(stack.pop_back()) return sorted_list def print_list(list: List[Int]): print('List','[', sep=': ', end=' ') for i in range(len(list)): print(list[i],end=' ') print("]", end='\n') my_list = List[Int]() for i in range(5): my_list.append(int(random_si64(0,100))) print("Original") print_list(my_list) sorted_list = stack_sort(my_list) print("Sorted") print_list(sorted_list)

Output:

Output
Original List: [ 14 97 69 34 80 ] Sorted List: [ 14 34 69 80 97 ]

In the above example, you can see that working with Mojo Lists are easy and very similar to Python Lists. The List functions I use are:

  1. Initializations: List[Int]() and List[Int](list). The former creates an empty List and the latter creates a deep-copy of an existing list 
  2. Append and pop_back: append function introduces a new value to the List, and pop_back removes the last element from the list.

See the List doc page for all the supported List methods. Lists are not confined to Int types, they can hold any Mojo type that implements the CollectionElement trait. Here is an example of a list holding Tensors of different sizes.

Mojo
from tensor import Tensor, rand t1 = rand[DType.float32](2,2) t2 = rand[DType.float32](3,2) t3 = rand[DType.float32](1) lst = List[Tensor[DType.float32]](t1,t2,t3) for i in range(len(lst)): print(lst[i])

Output:

Output
Tensor([[0.26290616393089294, 0.089547768235206604], [0.58222967386245728, 0.59191876649856567]], dtype=float32, shape=2x2) Tensor([[0.87663388252258301, 0.72621172666549683], [0.29710233211517334, 0.89949768781661987], [0.90153425931930542, 0.16471293568611145]], dtype=float32, shape=3x2) Tensor([[0.90684455633163452]], dtype=float32, shape=1)

Print supports sep and end keywords

Lists don’t support print() today, and I will use this as an opportunity to show you the new print() features. In the print_list() function above you can see that we use Python style sep and end keywords. The sep keyword argument specifies the string that should be used to separate the items that are passed to the print() function. The end keyword argument specifies the string that should be printed at the end of the print() function call. Using both, I wrote the list printing function in the previous example.

Mojo
def print_list(list: List[Int]): print('List','[', sep=': ', end=' ') for i in range(len(list)): print(list[i],end=' ') print("]", end='\n')

The output of calling our custom print_list():

Output
List: [ 14 97 69 34 80 ]

New ulp() function in math module

Moving on, there is a new addition to the math module in the standard library - the ulp() function. The ulp() function is also called Units of Least Precision, and it returns the distance between the floating-point number x and the next representable floating-point number in magnitude. This is useful to define a tolerance for comparing floating-point numbers, and to understand the precision of a floating-point number at a particular value. Let’s take a look at an example that uses ulp(). 

The Newton-Raphson method is a way to find a good approximation for the root of a function, i.e. the value of the variable that makes the function value zero. It’s an iterative algorithm that iterates on x using the following formula, until convergence is reached.

$$x_{n+1} = x_n - \frac{f(x_n)}{f'(x_n)}$$

We can break from the iterations if changes in x falls below some threshold which we define using math.ulp(). Consider this polynomial:

$$f(x) = x^3 - 2x^2 + 3x - 6$$

Now plug in the value $x=2$ and $f(x)$ become zero. The polynomial of degree 3 must have 2 other roots, and in this case they are imaginary, so we’ll ignore them for now. Here's Newton-Raphson's method of finding the root.

Mojo
from math.polynomial import polynomial_evaluate from math import abs, max, ulp alias dtype = DType.float64 alias width = simdwidthof[dtype]() alias coeff = List[Scalar[dtype]](-6, 3, -2, 1) alias deriv_coeff = List[Scalar[dtype]](3, -4, 3) def f(x: Float64) -> Float64: return polynomial_evaluate[dtype, 1, coeff](x) def df(x: Float64) -> Float64: return polynomial_evaluate[dtype, 1, deriv_coeff](x) def newton_raphson(x0: Float64, tolerance: Float64) -> Float64: x = x0 steps = 0 while True: x_new = x - f(x)/df(x) steps += 1 if abs(x_new - x) < max(tolerance, ulp(x_new)): print("Steps to convergence:", steps) break x = x_new return x_new initial_guess = 100 # Initial guess for the root tolerance = 1e-2 # Tolerance for terminating the algorithm root = newton_raphson(initial_guess, tolerance) print("Root found using Newton-Raphson method:", root) print("Value of the polynomial at the root:", polynomial_evaluate[dtype, 1, coeff](root))

In the above code we provide both an early stop tolerance value, and take the maximum of tolerance and (tolerance, ulp(x_new)) to decide if we should terminate our iterative algorithm.

If you run the code you’ll see that the algorithm finds the root of the polynomial and terminates after 15 steps.

Output
Steps to convergence: 15 Root found using Newton-Raphson method: 2.0 Value of the polynomial at the root: 0.0

If I want it to converge sooner, I can increase the manual tolerance to say 1e-2 and it converges in 13 steps with an approximate solutions

Mojo
initial_guess = 100 # Initial guess for the root tolerance = 1e-2 # Tolerance for terminating the algorithm root = newton_raphson(initial_guess, tolerance) print("Root found using Newton-Raphson method:", root) print("Value of the polynomial at the root:", polynomial_evaluate[dtype, 1, coeff](root))

Output:

Output
Steps to convergence: 13 Root found using Newton-Raphson method: 2.000010918736431 Value of the polynomial at the root: 7.6431631893613403e-05

math.ulp is useful to compare floating point numbers, and are especially useful in iterative algorithms used in machine learning and numerical optimization problems.

New Mojo <-> Python interoperability features

One of the highlights of this release for me is the ability to pass keyword arguments to Python functions. For example, in earlier versions if you wanted to use MatplotLib to create a figure you’d have to do it this way:

Mojo
fig = plt.figure() fig.set_size_inches(10, 6)

Now, with keyword arguments you can:

Mojo
plt.figure(figsize=(10, 6))

This makes using Python modules much easier. Building on our previous example of finding the root of the polynomial $f(x) = x^3 - 2x^2 + 3x - 6$, let’s visualize it using MatplotLib to see where the graph cross line y=0 and you can see that it crosses at x=2 which agrees with the value we estimated using Newton-Raphson’s method.

Here’s the Mojo code with Python interoperability and the new Python keyword arguments feature, that you can use to generate the graph.

Mojo
from python import Python from math.polynomial import polynomial_evaluate np = Python.import_module("numpy") plt = Python.import_module("matplotlib.pyplot") alias dtype = DType.float64 alias coeff = List[Scalar[dtype]](-6, 3, -2, 1) alias deriv_coeff = List[Scalar[dtype]](3, -4, 3) x_vals = np.linspace(-2, 3, 100) y_vals = PythonObject([]) y_deriv_vals = PythonObject([]) for i in range(len(x_vals)): y_vals.append(polynomial_evaluate[dtype, 1, coeff](x_vals[i].to_float64())) y_deriv_vals.append(polynomial_evaluate[dtype, 1, deriv_coeff](x_vals[i].to_float64())) plt.figure(figsize=(10, 6)) plt.plot(x_vals, y_vals, label="Polynomial:x^3 - 2x^2 + 3x - 6") plt.plot(x_vals, y_deriv_vals, label='Derivative of the Polynomial') plt.xlabel('x') plt.ylabel('y') plt.title("Plot showing root of x^3 - 2x^2 + 3x - 6") plt.legend() plt.grid(True) plt.axhline(0, color='black',linewidth=0.5) plt.axvline(0, color='black',linewidth=0.5) plt.axvline(x=2, color='r', linestyle='--', linewidth=2) plt.plot(2, 0, marker='o', markersize=15, color='b') plt.text(2.1, -2, 'Root', fontsize=12, color='b') plt.show()

New core language features

Mojo 24.2 also includes new core language features that makes it easier to write structs and to work with keyword arguments to functions.

Structs implicitly conform to traits

This is my favorite new core language feature, because I no longer have to specify all the traits in the Struct’s definition. For example, previous we had to define a Struct like this:

Mojo
from random import rand struct MojoArray[dtype: DType = DType.float64](Stringable, Sized): ... fn __len__(self) -> Int: ... fn __str__(self) -> String: ...

With this new feature, you can skip mentioning the traits, as long as you’ve implemented the methods in the Struct. Here is a working example:

Mojo
from random import rand struct MojoArray[dtype: DType = DType.float64](): var _ptr: DTypePointer[dtype] var numel: Int fn __init__(inout self, numel: Int, *data: Scalar[dtype]): self.numel = len(data) self._ptr = DTypePointer[dtype].alloc(len(data)) for i in range(len(data)): self._ptr[i] = data[i] fn __len__(self) -> Int: return self.numel fn __str__(self) -> String: var s: String = "" s += "[" for i in range(len(self)): if i>0: s+=" " s+=self._ptr[i] s+="]" return s var arr = MojoArray[DType.index](1,2,3,4,5,6,7) print("Array and length:") print(arr) print(len(arr))

Output

Output
Array and length: [2 3 4 5 6 7] Length: 6

However, as the changelog notes, we still strongly encourage you to explicitly list the traits a struct conforms to when possible.

Python style variadic keyword arguments: **kwargs

Mojo now has support for variadic keyword arguments: **kwargs to support an arbitrary number of keyword arguments. For example, here is a function that evaluates polynomials of user specified degree via coefficients provided as keyword arguments.

Mojo
from math.polynomial import polynomial_evaluate alias dtype = DType.float64 def polynomial(x: Float64, **kwargs: Int): val = 0.0 idx = 0 var s: String = "" for key in kwargs.keys(): if idx == 0: s += key[]+" + " else: s+= key[]+"*x^"+str(idx)+" + " val += kwargs[key[]]*(x**idx) idx+=1 print("Evaluating polynomial: [",s[:-2],'] ',"at",x) print(val)

Let’s evaluate a degree 3 polynomial function at x=42 and compare its results with the standard library’s polynomial_evaluate() function.

Mojo
polynomial(42.0, a=-6, b=3, c=-2, d=1) alias coeff = List[Scalar[dtype]](-6, 3, -2, 1) print("Compare results with math.polynomial") print(polynomial_evaluate[dtype, 1, coeff](42.0))

Output:

Output
Evaluating polynomial: [ a + b*x^1 + c*x^2 + d*x^3 ] at 42.0 70680.0 Compare results with math.polynomial 70680.0

Using keyword arguments, we can increase the degree of the polynomial by providing additional coefficients and compare the results:

Mojo
polynomial(42.0, a=-6, b=3, c=-2, d=1, e=2, f=3) alias coeff = List[Scalar[dtype]](-6, 3, -2, 1, 2, 3) print("Compare results with math.polynomial") print(polynomial_evaluate[dtype, 1, coeff](42.0))

Output

Output
Evaluating polynomial: [ a + b*x^1 + c*x^2 + d*x^3 + e*x^4 + f*x^5 ] at 42.0 398367768.0 Compare results with math.polynomial 398367768.0

New Mojo nightly build, nightly Visual Studio Code extension and open-source standard library

With the release of the open-source Mojo standard library, the Mojo public repo now has a new branch called Nightly, which is in sync with the Mojo nightly build. All pull requests should be submitted against the nightly branch, which represents the most recent nightly build. 

The nightly build also has an accompanying nightly Visual Studio Code extension. Note that you can only use one Mojo extension at a time. If you’re using the nightly extension, please disable the release extension.

Read more about contributing to open-source in our contributing guide on GitHub.

But wait, there is more!

We’re excited for you to try out the latest release of Mojo. For a detailed list of what’s new, changed, moved, renamed, and fixed, check out the changelog in the documentation. 

All the examples I used in this blog post are available in a Jupyter Notebook on GitHub, check it out!

https://github.com/modularml/devrel-extras/blob/main/blogs/whats_new_mojo_24_2

Additional resources to get started:

Until next time! 🔥

Next blog post:

View All Blogs

Developer

Shashank Prasanna
,
AI Developer Advocate
Shashank Prasanna
,
AI Developer Advocate

Shashank Prasanna

AI Developer Advocate

Shashank is an engineer, educator and doodler. He writes and talks about machine learning, specialized machine learning hardware (AI Accelerators) and AI Infrastructure in the cloud. He previously worked at Meta, AWS, NVIDIA, MathWorks (MATLAB) and Oracle in developer relations and marketing, product management, and software development roles and hold an M.S. in electrical engineering.

View more articles from Shashank