A Quick Introduction to Mojo — a Superset of Python That Is Super Fast

Learn a new sibling language of Python for AI

There is a brand new language called Mojo that appeared in 2023 and became very eye-catching quickly. It was created by Chris Lattner, the original architect of the Swift programming language. Mojo is designed to be a superset of Python. The syntax is very similar to Python and thus very friendly to Python developers.

Mojo is a compiled language and is magnitudes faster than Python. It’s specifically optimized for ML/AI-related calculations and thus can be interesting and helpful to learn for common Python developers.

In this post, we will introduce the very fundamentals of Mojo and compare them with their counterparts in Python. I’m sure you will be surprised by its syntax and efficiency when compared with Python.

Why Mojo

Mojo is currently a proprietary programming language developed by the AI company Modular, whose goal is to build an extensible and unified AI platform that removes complexity and adds flexibility and speed to AI development. And Mojo is the new language developed for this platform to solve the AI infrastructure scaling and acceleration issues. It is the first programming language built with MLIR, a compiler infrastructure that’s ideal for heterogeneous hardware including modern CPUs and GPUs. In a nutshell, Mojo is a new language for AI.

Mojo and Python

Mojo is designed to be a superset (or advanced version) of Python, similar to how TypeScript is a superset of JavaScript which is a great enhancement to JavaScript. Therefore, it will be very easy for Python developers to get started with Mojo.

Python is a dominant language in ML/AI and many other fields, especially in data-related ones. However, Python suffers from poor low-level performance and global interpreter lock (GIL) for concurrency. These issues have long been a headache issue for Python developers and are the embarrassing reason why people say Python is slow and not concurrent. They are also an important fact that limits the development of Python and what makes Python end up as a glue language where all core functionalities are developed in other more efficient compiled languages like C.

Mojo is designed to fill the gap between the popularity and simplicity of Python and the efficiency of compiled languages like C and Rust. It actively embraces the Python community and tries to make the language easy to learn and use by hiding all the complexity like compilation and memory management for efficient execution of the code behind the scenes.

In this post, we will guide you through the basics of Mojo by making direct comparisons with Python. We will introduce the language basics like variables, functions, and structs. Finally, we will compare the efficiency of Mojo and Python for the same piece of code and how to run Python code directly in Mojo. You will then get a better picture of this new language and can see the potential of usage in your specific area.

Install Mojo

We can install Mojo on a Linux computer with these commands:

curl https://get.modular.com | sh - && \
modular auth mut_73b76eabd7a04555be4daa751d3e7088

modular install mojo

The instructions for other platforms can be found here.

When the installation is completed, you will see the instructions to set up Mojo in the console:

Copy and run the corresponding commands for your shell and you will be able to run Mojo:

The usage of the print command is the same as in Python, and we have just written the first “Hello World” code in Mojo!

Note that to exit Mojo REPL, we should type a colon : plus quit, not quit() as in Python, or exit as in iPython.

Variables

Similar to the case in Python, we can create a variable in Mojo with just a name and value:

  1> x = 1 
  2.  
(Int) x = 1

A variable declared like this is mutable, meaning the value can be changed:

  2> x = 2 
  3. print(x) 
  4.  
2

However, since Mojo is strictly typed, you cannot change the type of a variable:

  4> x = "a" 
  5.  
[User] error: Expression [2]:1:5: cannot implicitly convert 'StringLiteral' value to 'Int' in assignment
x = "a"

Declaring a variable with just a name and value is to be consistent with Python, but it’s not seen as a good practice in Mojo. Instead, the var and let keywords should be used:

var is used to define mutable variables.
let is used to define immutable variables.

In the example shown above, a variable declared without var or let is mutable as it defaults to var.

If we attempt to change the value of a variable defined by let, an error would be raised:

  1> let x = 1 
  2. x = 2 
  3. print(x) 
  4.  
[User] error: Expression [0]:2:1: expression must be mutable in assignment
x = 2

Although Mojo is designed to be a “superset” of Python, it currently is not yet. Therefore, it just looks like Python, but the syntax can be different sometimes. For example, the keywords for type annotations are different:

  1> var x: Int = 1 
  2. var y: String = "a" 
  3. var z: Bool = True

Functions

Mojo functions can be declared with the fn or def keyword.

The fn declaration enforces type-checking and memory-safe behaviors (Rust style), while def allows no type declarations and dynamic behaviors (Python style).

When we use def to declare a function, we don’t need to specify the argument types and the return type and can create variables with just a name and value, exactly as what we can do in Python:

def sum_up(n):
    sum = 0
    for i in range(1, n+1):
        sum += i
    
    return sum

The same function defined with the fn keyword would be:

fn sum_up(n: Int) -> Int:
    var sum: Int = 0
    for i in range(1, n+1):
        sum += i
    
    return sum

The function defined with fn is strictly typed and provides compile-time checks to ensure the function receives and returns the correct types. Supporting the def keyword is just to make it easier for Python users to understand the syntax of Mojo and get started with it. It is not seen as a good practice when writing Mojo code. When you read more Mojo code, you will see almost all functions are defined with the fn keyword.

Mojo Structs

A struct in Mojo is similar to the class in Python. They both support properties, methods, decorators, etc.

Let’s define a simple class in Python:

class Dog:
    def __init__(self, name):
        self.name = name

    def run(self):
        print(f"{self.name} is running.")

dog = Dog('Teddy')
dog.run()

The corresponding struct in Mojo is:

struct Dog:
    var name: String

    fn __init__(inout self, name: String):
        self.name = name

    fn run(self):
        print(self.name, " is running.")

let dog = Dog('Teddy')
dog.run()

Some key differences here:

We need to define the properties first and then assign values in the constructor (__init__()).
The first argument to __init__() is called self as a convention as in Python, and it’s defined with the keyword inout, which means self is a mutable reference. It’s a complex concept but we don’t need to dive that deep in this post.
f-string is not supported in Mojo yet and thus we need to use regular printing.

Once you know the basics of variables and functions, it’s easy to learn the syntax of structs which can be seen as the combination of the two. However, structs in Mojo is a complex concept and covers everything in Mojo if you want to dive deeper. Even basic types like Int, String, and Bool are structs behind the scenes, and that’s why the first characters are capitalized like classes/structs.

Compare the efficiency of Python and Mojo code

OK, now that we know the very basic syntax of Mojo, let’s compare the efficiency of Python and Mojo code and see why it’s amazing and has attracted so much interest recently.

Let’s put the following code in a Python file named sum_up.py:

def sum_up(n):
    sum = 0
    for i in range(1, n+1):
        sum += i
    
    return sum

def main():
    print(sum_up(1000_000_000))

if __name__ == "__main__":
    main()

And put this code in a Mojo file named sum_up.mojo:

fn sum_up(n: Int) -> Int:
    var sum: Int = 0
    for i in range(1, n+1):
        sum += i
    
    return sum

fn main():
    print(sum_up(1000_000_000))

Note that for the Mojo file, we must define a function named main() as the entry point, but we don’t need to call it.

In Linux, we can use the time command to measure the execution time of a command. Let’s run the Mojo file first, otherwise, you may lose your patience if we run the Python file first 😅:

$ time mojo sum_up.mojo 
500000000500000000

real    0m0,134s
user    0m0,161s
sys     0m0,004s

The code was completed instantly, in less than a second!

Now let’s try with the same code written with our old friend Python:

$ time python sum_up.py 
500000000500000000

real    0m41,312s
user    0m41,288s
sys     0m0,007s

This time it took more than 40 seconds. More than 100 times slower!

The root reason is that Mojo is a compiled language and Python is an interpreted one. It’s not surprising that the same code in a compiled language can be 100 times faster than in Python. However, it’s pretty surprising that the syntax is so similar to Python and thus so friendly to Python users. If all our data processing and ML/AI code could be 100 times faster with minimal changes to our code, it would be fantastic.

Run Python code directly in Mojo

In the long run, when Mojo is mature and becomes a real superset of Python, we should be able to use Mojo immediately and run Python code directly with Mojo, just like we can write valid JavaScript code in TypeScript.

However, currently, we can only import Python modules in Mojo and then call Python functions and interact with Python objects from Mojo code. Let’s create a new file called call_python_in_mojo.mojo and import sum_up.py module created above and run it in Mojo directly:

from python import Python

fn main() raises:
    Python.add_to_path(".")
    let mypython = Python.import_module("sum_up")

    let result = mypython.sum_up(1000_000_000)
    print(result)

Note that we need to add the path to the Python module so it can be found by Mojo. We can run this code with this command:

time mojo call_python_in_mojo.mojo

You will find out the speed is the same as running the code with Python directly, which is a bit disappointing. However, it should be noted that Mojo is still in early development and many features are not yet implemented. It is expected that the usability and integration with Python will be much better over time.

In this post, we have introduced the very basics of Mojo, a new programming language that is designed to be a superset of Python and thus has very similar syntax to Python. We have experienced the simple syntax and the amazing speed of Mojo code which can be very attractive for data processing and machine learning.

However, Mojo is still very new and not quite mature yet. It still needs some time to be ready for usage by common developers in their work. Nonetheless, it’s worth the time to keep track of this language and don’t fall behind in the new era of AI. And when it’s ready, we can quickly switch to boost our efficiency, which can potentially be a game changer event.

If you enjoyed this article, consider trying out the AI service I recommend. It provides the same performance and functions to ChatGPT Plus(GPT-4) but more cost-effective, at just $6/month (Special offer for $1/month). Click here to try ZAI.chat.

Summarize