Mojo Programming for AI: Better than Python?

If you’re passionate about staying at the cutting edge of AI innovation, you’re in the right place! As we lunge into the future, being at the forefront of AI technology is no longer just an advantage—it’s a necessity. This is where Mojo programming comes into the fray.

Now, most of us rightfully associate AI / ML programming with Python. So where does Mojo programming language come in? That’s what this follow-up article on Mojo aims to find out after we explored the fundamentals of the Mojo language in our previous article.

Why Choose Mojo Programming Language?

You wouldn’t be silly to ask this question when Mojo programming language hasn’t even publically released yet. In fact, it’s a smart question. But do you know what it means? Opportunity. And do you know who is behind this new language? Chris Lattner, the creator of Swift programming language. We know how that turned out. And he’s also the co-founder of LLVM, Clang compiler, and MLIR compiler. So, this is the best time to fine-tune your skills to specialize in a language most people will not know about until later.

And choosing Mojo programming language for AI development carries significant benefits. Modular’s powerful AI engine, Mojo’s ability to support existing Python code and infrastructure, and its high performance make it a compelling choice for developers. That’s not all! Mojo language also aims to be a complete superset of Python before release.

But is the hype real, especially when it comes to its performance?

Mojo Programming Language, Analyzed

While Python has been the go-to language for various tasks over the years, Mojo’s features are specifically tailored for AI and systems development. Remember that Mojo programming language aims to deal with the pain points.

For instance, it offers in-built support for concurrent and distributed computing and parametric metaprogramming—fundamental attributes for today’s data-intensive AI applications. This makes Mojo a strong contender compared to Python in AI development.

And yet… how exactly does it compare to Python? Is Mojo is better than Python? Let’s take a deeper look.

Python vs Mojo: A Comprehensive Comparison

What’s really exciting is that the Mojo playground allows you to use the default Python interpreter for the Python code versions of any program. It will make more sense once I show you how.

All you have to do is add this to the top of your Jupyter notebook cell in the Mojo programming playground to indicate that the cell contains Python code.:

%%python

Then, you can write any Python code, and it will run as if on native Python.

For instance, this compiles in the Mojo playground even though it doesn’t natively support lists or list functions. It works because it executes using the Python interpreter.

%%python
data = [1, 2, 3, 4, 5]

# Calculate the mean (average) of the numbers
mean = sum(data) / len(data)

# Print the mean
print(mean)

#output: 3.0

Note that the variables, functions, and imports defined in a Python cell are available for access in future Mojo cells. This lets you write and execute normal Python code within the Mojo programming playground. And you can neatly integrate it with native Mojo language code too. How futuristic is that?

That said, it’s time to explore the key differences between Python and Mojo programming language.

Basic Data Types: How are Integer and String Different in Mojo Programming Language

If you have been paying attention, you must have noticed that Mojo uses Int (with strong type checking), with a capital I, compared to int in Python, with a lowercase i. This is on purpose.

Mojo’s Int type aims to be simple, fast, and perform better. On the other hand, int in Python is dynamically allocated space and doesn’t have a fixed limit on its size. Mojo’s Int is a fixed 32-bit signed integer, while Python’s int is an arbitrary precision integer. Of course, a fixed-size type also presents problems like overflow and underflow, which are not present in Python.

Mojo fixes the slower performance of Python’s data types by utilizing stronger type checking, like in C++ and Swift (Lattner, remember?). Thus, Mojo’s Int will likely have more consistent and faster performance compared to Python’s int. Especially when it comes to larger data sets or more complex calculations.

Here’s an example highlighting the difference in the variable declaration for integer type.

#mojo encourages strong type checking but allows implicit variable data types too
var number: Int = 10
#this also works:
var number = 10

#python does not support explicit or strong type checking
number = 10

In Python, the variable dynamically takes on the type for the value of 10. And the string data type is also different. In Mojo programming, the String type is imported from the String module. Instead of using the str() method for type conversion as in Python, Mojo uses the String() constructor. Another key difference is that Mojo only supports commas to separate or concatenate strings and variables.

To work with strings, you must add this line to your code:

from String import String

#Strings in Mojo langauge, including the import, String constructor, and ',' for concat

from String import String
let name: String = "Umar"
let age: Int = 99
print("Hello, my name is ", name, "and I'm ", String(age), "years old.")

# Strings in Python using built-in str data type, str() method for type conversion, and supports '+' or ',' for concat

name = "UmarB"
age = 92
print("Hello, my name is " + name + " and I'm " + str(age) + " years old.")  # Output: Hello, my name is Alice and I'm 25 years old.

The outputs of both code snippets are similar, as shown below, except for better spacing in the Python version.

Beyond data types, the differences also extend to function declaration.

Comparing ‘fn’ and ‘def’ Functions in Mojo Programming: Unveiling Control, Mutability, and Scoping Differences

Both fn and def keywords can be used to define functions in Mojo, are interchangeable at the interface level, and have parameters and return values. They can also raise exceptions using the raises function effect.

All values passed in Mojo functions use value semantics by default, unlike Python. Mojo functions, by default, receive a copy of arguments and can change them inside the function scope with no visible change outside.

Here is a list of comparison between fn and def:

• fn provides a more controlled and restricted environment inside its body compared to def.
• In a fn declaration, argument values default to being immutable and require explicit type specifications, similar to let declarations, while in def they default to being mutable and argument declaration doesn’t have to be explicit.
• The inout modifier in fn indicates that the parameter can be mutated, and the changes will be reflected outside the function, like in Python. def functions do not support the inout parameter modifier, while fn functions allow specifying inout parameters.
• fn functions require explicit type annotations for parameters and the return type. If the return type is not explicitly declared, it is interpreted as returning None. In contrast, def functions can have implicit return types, where if a return statement is missing, it will implicitly return None.
• fn functions support the owned parameter modifier, which indicates ownership transfer, allowing explicit control over memory management, while def functions do not have an equivalent modifier.

from String import String

def greet(name):
    print("Hello,")
    print(name)

fn farewell(name: String):
    print("Goodbye, ", name)

# Both can be called with the same syntax
greet("Alice")
farewell("Bob")

def add(x, y):
    return x + y

fn subtract(x: Int, y: Int) -> Int:
    return x - y

print(add(1,2))
print("Subtract: ", subtract(5,2))

Here is the output of the code above:

To recap, `def` is defined by necessity to be very dynamic, flexible and generally compatible with Python: arguments are mutable, local variables are implicitly declared on first use, and scoping isn’t enforced. This is great for high level programming and scripting, but is not always great for systems programming. To complement this, Mojo provides an `fn` declaration which is like a “strict mode” for `def`.

Note:

It’s not possible to use print(“Hello, “, name) inside the def function. This is the error you get: no matching function in call to ‘print’:

Struct vs. Class: Contrasting Data Containers in Mojo and Object-Oriented Classes in Python

The Mojo programming language does not support classes yet. Currently, it only supports structures. Structs in Mojo programming, similar to classes in Python language, can have member variables and methods. They are also static and inline their data within their container. On the contrary, classes in Python have a more dynamic nature and provide extensive support for OOP concepts.

Here is a comprehensive comparison between structs and classes:

• Structs provide a static and compile-time bound approach, unlike the dynamic nature of classes in Python.
• Structs in Mojo are inlined into their container instead of being implicitly indirect and reference counted. Structs primarily focus on data representation rather than behavior. Classes, on the other hand, allow for inheritance and polymorphism, facilitating code reuse and extensibility.
• Structs offer indirection-free field access, allowing direct access to fields without method calls. But classes can have methods that encapsulate behavior and can be called on instances of the class.
• Structs are often used for performance-critical scenarios and low-level systems programming. However, classes commonly building complex software systems, model real-world entities, and organize code in a modular way.

#creating a Person struct with functions to greet

from String import String
@value
struct Person:
    var name: String
    var age: Int

    fn greet(self):
        print("Hello, my name is ", self.name, "and I'm ", String(self.age), "years old.")

# Creating an instance of the Person struct
person = Person("Umar", 22)
person.greet()


#creating a class and methods in Python for greeting a person

%%python
class Person:
    def __init__(self, name, age):
        self.name = name
        self.age = age

    def greet(self):
        print("Hello, my name is " + self.name + " and I'm " + str(self.age) + " years old.")

# Creating an instance of the Person class
person = Person("Bukari", 30)
person.greet()

Note that as I pointed out in the last article, a complete guide to Mojo language, using the @value decorator allows you to skip boilerplate code such as the init method. As shown in the code above, the class uses an init method. However, in the struct in Mojo uses the decorator to create the boilerplate methods implicitly.

This is the code output in Mojo programming:

Importing Python Libraries in Mojo Programming: Leveraging Python’s Existing Rich Ecosystem

The cool thing about Mojo programming language is that you can use Python’s existing rich and versatile ecosystem of libraries! Don’t believe me? Give me a moment, and I will show you.

At the top of your Mojo programming code, add this line:

from PythonInterface import Python

After that, you can use the import_module function to add any Python library to your Mojo program.

from PythonInterface import Python
# This is equivalent to Python's `import numpy as np`
let np = Python.import_module("numpy")


# Now use numpy as if writing in Python
array = np.array([1, 2, 3])
print(array)

Just like that, you can use a variety of other Python libraries already!

And now, let’s discuss the fabled high performance of Mojo. Is it really so fast?

Performance Comparison: Mojo vs Python language

The CEO of Modular, Chris Lattner claims Mojo is 35 THOUSAND times faster than Python. Yes, you read that right. Can you even imagine that? Note that this metric is is for Mojo-optimized code only.

So, for now, we will test a matrix multiplication Python function and then convert the code to Mojo language. After that, we will compare the benchmark for performance simply by using the same Python implementation in Mojo code.

Here is the Python implementation of matrix multiplication on the Mojo playground:

%%python
import numpy as np
from timeit import timeit

class Matrix:
    def __init__(self, value, rows, cols):
        self.value = value
        self.rows = rows
        self.cols = cols
        
    def __getitem__(self, idxs):
        return self.value[idxs[0]][idxs[1]]
    
    def __setitem__(self, idxs, value):
        self.value[idxs[0]][idxs[1]] = value

def benchmark_matmul_python(M, N, K):
    A = Matrix(list(np.random.rand(M, K)), M, K)
    B = Matrix(list(np.random.rand(K, N)), K, N)
    C = Matrix(list(np.zeros((M, N))), M, N)
    secs = timeit(lambda: matmul_python(C, A, B), number=2)/2
    gflops = ((2*M*N*K)/secs) / 1e9
    print(gflops, "GFLOP/s")
    return gflops

def matmul_python(C, A, B):
    for m in range(C.rows):
        for k in range(A.cols):
            for n in range(C.cols):
                C[m, n] += A[m, k] * B[k, n]

After that, run this command in the next cell:

python_gflops = benchmark_matmul_python(128, 128, 128).to_float64()

This is the output on my machine (MacBook M1):

Next, let’s modify this Python code for a Mojo program implementation:

#|code-fold: true
#|code-summary: "Import utilities and define `Matrix` (click to show/hide)"

from Benchmark import Benchmark
from DType import DType
from Intrinsics import strided_load
from List import VariadicList
from Math import div_ceil, min
from Memory import memset_zero
from Object import object, Attr
from Pointer import DTypePointer
from Random import rand, random_float64
from TargetInfo import dtype_sizeof, dtype_simd_width

# This exactly the same Python implementation, 
# but is infact Mojo code!
def matmul_untyped(C, A, B):
    for m in range(C.rows):
        for k in range(A.cols):
            for n in range(C.cols):
                C[m, n] += A[m, k] * B[k, n]
fn matrix_getitem(self: object, i: object) raises -> object:
    return self.value[i]


fn matrix_setitem(self: object, i: object, value: object) raises -> object:
    self.value[i] = value
    return None


fn matrix_append(self: object, value: object) raises -> object:
    self.value.append(value)
    return None


fn matrix_init(rows: Int, cols: Int) raises -> object:
    let value = object([])
    return object(
        Attr("value", value), Attr("__getitem__", matrix_getitem), Attr("__setitem__", matrix_setitem), 
        Attr("rows", rows), Attr("cols", cols), Attr("append", matrix_append),
    )

def benchmark_matmul_untyped(M: Int, N: Int, K: Int, python_gflops: Float64):
    C = matrix_init(M, N)
    A = matrix_init(M, K)
    B = matrix_init(K, N)
    for i in range(M):
        c_row = object([])
        b_row = object([])
        a_row = object([])
        for j in range(N):
            c_row.append(0.0)
            b_row.append(random_float64(-5, 5))
            a_row.append(random_float64(-5, 5))
        C.append(c_row)
        B.append(b_row)
        A.append(a_row)

    @parameter
    fn test_fn():
        try:
            _ = matmul_untyped(C, A, B)
        except:
            pass

    let secs = Float64(Benchmark().run[test_fn]()) / 1_000_000_000
    _ = (A, B, C)
    let gflops = ((2*M*N*K)/secs) / 1e9
    let speedup : Float64 = gflops / python_gflops
    print(gflops, "GFLOP/s, a", speedup.value, "x speedup over Python")
    

In the next cell, run this line:

benchmark_matmul_untyped(128, 128, 128, 0.0022926400430998525)

On my machine, this code runs nearly 5x times faster, with minimal optimization. We simply converted the Python code to Mojo’s native code:

I also want to show you this Mojo code from the playground:

The optimized Mojo code for matrix multiplication runs 14000x times faster. Even if it’s not the same on your machine, it is still unprecedented!

Aren’t you interested to learn more about Mojo programming’s role in AI development now?

Mojo’s Role in AI Development: The Need for Powerful Tools in a Competitive Future

Modular‘s AI engine natively supports dynamic shapes for AI workloads, based on a new technology called Shapeless, outpacing other statically shaped compilers. Shapeless allows Modular to represent and manipulate variable-length inputs without having to know their exact shapes in advance. The AI engine is also fully compatible with existing frameworks and servers, and it allows developers to write their own custom operators. As a result, Modular’s engine can deploy AI models in a variety of environments, including the cloud, on-premises, and edge devices.

Impressively, the Modular AI Engine exhibits speedups of 3x-9x versus the default TensorFlow on BERT, a highly optimized language model. While XLA necessitates a significantly longer compile time than TensorFlow, it does offer superior execution performance. Despite this, the Modular AI Engine continuously outperforms, delivering a 2x-4x faster and superior performance than XLA.

Comparing Mojo’s AI Engine to Python Engines:

• Mojo’s AI engine is a good choice if you need to build a model with dynamic shapes. Python engines can also build models with dynamic shapes, but they may require more code and be less efficient.
• If you need to deploy your model on various platforms, Mojo’s AI engine is a good choice. Python engines can also be deployed on a variety of platforms, but Mojo’s AI engine is designed to be more portable.
• If you need to write custom operators for your model, Mojo’s AI engine is a good choice. Python engines can also write custom operators, but Mojo’s AI engine makes it easier.

Conclusion: The Future Scope of Mojo Programming for AI

When you choose to embark on the journey of Mojo programming for AI development, you’re not only adopting a programming language—you’re aligning with a philosophy that prioritizes performance, flexibility, and user-friendliness, as explored in our recent article: Mojo Language 101. As Mojo progresses to be a complete superset of Python before its public release, early adoption and specialization in Mojo programming could offer significant advantages in the future.

The emerging Mojo programming language’s tailored features for AI and systems development, along with compatibility with Python code and libraries, make it a potential game-changer. With distinct advantages over Python, including high performance, in-built support for concurrent and distributed computing, and superior handling of data types, Mojo positions itself as a powerful tool for AI development.

From matrix multiplication benchmarks that indicate Mojo to be up to five times faster than Python, to the staggering claim of fully optimized Mojo code running 14,000 times faster, the future of AI and systems development could be revolutionized. But Mojo’s proficiency is not limited to its speed; it also boasts superior compatibility and a more streamlined approach to handling dynamic shapes in AI workloads!

Please comment your thoughts on Mojo programming language and share with your friends if you found it helpful.

Related articles: Mojo Language Complete Guide and Classification of NSFW Images.

Written by: Syed Umar Bukhari.