Set up​

Let's start by creating a Python project and installing the necessary tools.

This will install the modular package along with PyTorch (torch), PIL for image processing (pillow), and NumPy for array operations. When you install the modular package, you'll get access to the max Python APIs and the Mojo compiler—everything needed to build high-performance custom operations.

Build the PyTorch interface​

Let's start by creating the PyTorch side of our integration. We'll build a simple function that uses our custom operation, but first we need to establish the interface.

Create a new file called grayscale.py and add the following code:

from std.pathlib import Path
import torch

@torch.compile
def grayscale(pic: torch.Tensor):
    output = pic.new_empty(pic.shape[:-1])  # Remove color channel dimension
    # We'll call our custom operation here
    return output

The grayscale function transforms a color image (with red, green, and blue channels) into a single-channel grayscale image. While PyTorch has built-in operations for this, implementing it as a Mojo custom op demonstrates the integration pattern you can apply to any performance-critical computation. Now we need to bridge to our Mojo implementation.

Integrate the custom operation​

The max.torch module provides CustomOpLibrary, which allows you to load and use compiled Mojo operations directly in PyTorch.

Update the grayscale.py file to include the following code:

from std.pathlib import Path
import torch
from max.experimental.torch import CustomOpLibrary

# Load the compiled Mojo package containing our custom operations
mojo_kernels = Path(__file__).parent / "operations"
ops = CustomOpLibrary(mojo_kernels)

@torch.compile
def grayscale(pic: torch.Tensor):
    output = pic.new_empty(pic.shape[:-1])  # Remove color channel dimension
    ops.grayscale(output, pic)  # Call our Mojo custom op
    return output

The CustomOpLibrary loads operations from a precompiled Mojo package file (.mojoc or .mojopkg). Once loaded, you can call these operations just like any other PyTorch function. The operations automatically handle data movement between PyTorch and MAX, and integrate with PyTorch when needed.

The compilation of your Mojo code into a .mojoc file is handled when you run your Python script. You don't need to manually invoke the Mojo compiler or manage build steps.

Implement the Mojo kernel​

Now for the core implementation—our high-performance Mojo kernel. We'll create a Mojo package that defines our grayscale conversion operation

Create a new file called grayscale.mojo inside the operations folder and add the following code:

from compiler import register
from max.tensor import InputTensor, OutputTensor, foreach
from std.gpu.host import DeviceContext

@register("grayscale")
struct Grayscale:
    @staticmethod
    def execute[
        target: StaticString,
    ](
        img_out: OutputTensor[dtype=DType.uint8, rank=2, ...],
        img_in: InputTensor[dtype=DType.uint8, rank=3, ...],
        ctx: DeviceContext,
    ) raises:
        @parameter
        @always_inline
        def color_to_grayscale[
            simd_width: Int
        ](idx: Coord) -> SIMD[DType.uint8, simd_width]:
            @parameter
            def load(idx: Coord) -> SIMD[DType.float32, simd_width]:
                return img_in.load[simd_width](idx).cast[DType.float32]()

            var row = idx[0]
            var col = idx[1]

            # Load RGB values
            var r = load((row, col, 0))
            var g = load((row, col, 1))
            var b = load((row, col, 2))

            # Apply standard grayscale conversion formula
            var gray = 0.21 * r + 0.71 * g + 0.07 * b
            return min(gray, 255).cast[DType.uint8]()

        foreach[color_to_grayscale, target=target, simd_width=1](img_out, ctx)

First, the @register("grayscale") decorator makes this operation available to PyTorch under the name grayscale. This is the ops.grayscale(output, pic) function called in the PyTorch model.

Then, the foreach primitive automatically parallelizes the operation across available compute units.

The color_to_grayscale function takes an index list and returns a SIMD vector of the grayscale value. The foreach primitive automatically parallelizes the operation across available compute units.

Finally, the target parameter allows the same code to run on both CPU and GPU.

This approach allows Mojo to handle the low-level optimization details: memory layout, vectorization, and parallelization, while you focus on the algorithm.

Run the example​

Now that we have a PyTorch model and a Mojo kernel, we can test it with a real image.

Create a new file called main.py to the root of your project and add the following code:

import io
import numpy as np
from PIL import Image
import torch
from grayscale import grayscale

# Specify a test image
image = Image.open("test_image.jpg")
# Convert to PyTorch tensor and move to GPU for processing
image_tensor = torch.from_dlpack(np.array(image)).cuda()

# Apply our custom grayscale operation
gray_image = grayscale(image_tensor)

# Convert back to PIL Image for display or further processing
result = Image.fromarray(gray_image.cpu().numpy())

You can run the example with the following command:

The following image shows the result of the grayscale operation.

This example demonstrates the complete pipeline: downloading an image, converting it to a PyTorch tensor, processing it with our Mojo custom op, and converting back to a standard Python image format. The operation automatically leverages GPU acceleration when available, providing significant performance improvements over CPU-only implementations.

Next steps​

Now that you wrote a Mojo kernel with your PyTorch model, check out these other tutorials:

• Build custom ops for GPUs: Write custom GPU and CPU operations in Mojo and load them into a MAX graph.
• Build a model graph with Module: Learn how to compose modules, write custom modules with explicit weights, load checkpoint data, and construct a model graph.