Skip to content

Interactive Mojo integration for Python notebooks - three patterns for running high-performance Mojo code from marimo

License

Notifications You must be signed in to change notification settings

DataBooth/mojo-marimo

Repository files navigation

mojo-marimo 🔥

Notebook-agnostic Mojo integration for Python - three patterns for running high-performance Mojo code from any Python environment.

Status:Beta - Three working approaches, evolving based on real-world usage

Note: Despite the name, this library works with any Python notebook system (Jupyter, marimo, VSCode, Google Colab) and even standalone Python scripts. The name reflects its origin, but a rename is under consideration—see ROADMAP.md.

Overview

mojo-marimo provides three distinct patterns for executing Mojo code from Python, each with different trade-offs:

  1. Decorator (@mojo) - Clean Pythonic syntax with template parameters and caching
  2. Executor (run_mojo()) - Dynamic execution from strings or files, great for code generation
  3. Extension Module - Compiled .so files for zero-overhead FFI calls (~1000× faster than subprocess)

Works everywhere: Jupyter notebooks, marimo, VSCode notebooks, Google Colab, IPython REPL, or standalone Python scripts. The core library has no notebook-specific dependencies.

from mojo_marimo import mojo

@mojo
def fibonacci(n: int) -> int:
    """
    fn fibonacci(n: Int) -> Int:
        if n <= 1:
            return n
        var prev: Int = 0
        var curr: Int = 1
        for _ in range(2, n + 1):
            var next_val = prev + curr
            prev = curr
            curr = next_val
        return curr
    
    fn main():
        print(fibonacci({{n}}))
    """
    pass

# Use like normal Python!
result = fibonacci(10)

Why This Exists

Python notebooks are brilliant for exploration, but hit a wall when you need serious performance. Traditional solutions—rewriting in C/C++, using Numba/JAX, or pure Python optimisation—are either too complex or don't deliver the 10-100× speedup you need.

The question: What if you could write high-performance Mojo code and run it interactively from Python notebooks with minimal friction?

This matters for:

  • Data scientists exploring algorithms that need real performance
  • Quant developers prototyping trading strategies or risk models
  • ML engineers benchmarking preprocessing pipelines or custom operators
  • Educators teaching performance engineering with immediate visual feedback

Features

Current (v0.1.0)

  • Three integration patterns (decorator, executor, extension modules)
  • Works with any Python environment (Jupyter, marimo, VSCode, IPython, scripts)
  • Interactive example notebooks in marimo and Jupyter (.ipynb) formats
  • SHA256-based binary caching (~/.mojo_cache/binaries/)
  • Pre-compilation validation (catches common syntax errors)
  • Cache management utilities (clear_cache(), cache_stats())
  • Monte Carlo and Mandelbrot examples with visualisation
  • 44 passing tests (75% coverage)
  • Comprehensive documentation + roadmap

Planned

See ROADMAP.md for full details:

  • Validate Jupyter compatibility with real-world testing
  • Auto-generate extension module boilerplate
  • Pattern library for common algorithms
  • Enhanced error handling and debugging
  • Multiple Mojo version support
  • Potential package rename (community feedback requested)

Installation

mojo-marimo supports both uv (recommended) and pixi for environment management.

Prerequisites

None! Mojo is now installed automatically as a Python package dependency.

Option 1: uv (Recommended)

# Clone the repository
git clone https://github.com/databooth/mojo-marimo
cd mojo-marimo

# Install dependencies (includes mojo)
uv sync --extra dev

# Verify setup
uv run python scripts/verify_setup.py

# Launch example notebook
uv run marimo edit notebooks/example_notebook.py

Option 2: pixi

# Clone the repository
git clone https://github.com/databooth/mojo-marimo
cd mojo-marimo

# Install with pixi
pixi install

# Verify setup
pixi run test-setup

# Launch example notebook
pixi run notebook-example

Quick Start

Pattern 1: Decorator (Recommended)

from mojo_marimo import mojo

@mojo
def sum_squares(n: int) -> int:
    """
    fn sum_squares(n: Int) -> Int:
        var total: Int = 0
        for i in range(1, n + 1):
            total += i * i
        return total
    
    fn main():
        print(sum_squares({{n}}))
    """
    pass

# Use like normal Python!
result = sum_squares(10)  # First call: ~1-2s, subsequent: ~10-50ms
print(result)  # 385

Pattern 2: Executor

from mojo_marimo import run_mojo

mojo_code = """
fn compute(n: Int) -> Int:
    return n * n

fn main():
    print(compute(42))
"""

result = run_mojo(mojo_code)  # Or run_mojo("path/to/file.mojo")
print(result)  # "1764"

Pattern 3: Extension Module

import mojo.importer  # Enables auto-compilation of .mojo → .so
import monte_carlo_mojo_ext

# Direct FFI call - no subprocess overhead!
x, y, inside, pi_est, error = monte_carlo_mojo_ext.generate_samples(1_000_000)
print(f"π ≈ {pi_est:.6f} ± {error:.6f}")

See examples/ and notebooks/ for complete working examples.

Usage

Verify Setup

# Using just
just test-setup

# Using uv
uv run python scripts/verify_setup.py

# Using pixi
pixi run test-setup

This checks that mojo is available and tests both approaches.

Interactive Notebooks

# Using just
just notebook-decorator  # @mojo decorator examples
just notebook-executor   # run_mojo() examples
just benchmark          # Performance comparison

# Using uv
uv run marimo edit notebooks/pattern_decorator.py
uv run marimo edit notebooks/pattern_executor.py
uv run marimo edit notebooks/benchmark.py

# Using pixi
pixi run notebook-decorator
pixi run notebook-executor
pixi run benchmark

Command-Line Demos

# Using just
just demo-examples
just demo-decorator

# Using uv
uv run python examples/examples.py
uv run python -m mojo_marimo.decorator

# Using pixi
pixi run demo-examples
pixi run demo-decorator

Performance Comparison

Testing on Apple Silicon (M-series) with fibonacci(10), sum_squares(100), and is_prime(104729):

Approach First Call Subsequent Calls Use Case
Uncached ~50-200ms ~50-200ms Development, debugging
Cached ~1-2s ~10-50ms Repeated execution
Decorator ~1-2s ~10-50ms Production, clean code

Key insights:

  1. Caching wins for repeated calls: 5-10× faster after first compilation
  2. Decorator has zero performance cost: Same speed as explicit caching, better developer experience
  3. All approaches deliver real Mojo performance: No Python fallbacks or compromises

See the benchmark notebook for detailed comparisons.

Documentation

Project Documentation

Blog & Announcements

Project Structure

mojo-marimo/
├── src/
│   └── mojo_marimo/          # Core library
│       ├── executor.py       # Cached Mojo execution
│       ├── decorator.py      # @mojo decorator
│       ├── validator.py      # Pre-compilation validation
│       └── __init__.py       # Package exports
├── examples/                 # Example implementations
│   ├── examples.py          # Python wrappers (fibonacci, etc.)
│   ├── examples.mojo        # Standalone Mojo code
│   └── reference/           # Reference .mojo files
├── benchmarks/              # Performance benchmarking
│   ├── python_baseline.py   # Pure Python implementations
│   ├── mojo_implementations.py  # Mojo implementations
│   ├── python_vs_mojo.py    # Python vs Mojo comparison notebook
│   └── execution_approaches.py  # Execution patterns comparison
├── notebooks/               # Interactive marimo notebooks
│   ├── pattern_decorator.py # @mojo decorator examples
│   ├── pattern_executor.py  # run_mojo() examples
│   └── interactive_learning.py  # Learning notebook
├── scripts/                 # Utility scripts
│   └── verify_setup.py      # Setup verification
├── tests/                   # Test suite
├── docs/                    # Documentation
│   ├── project/            # Project docs (contributing, changelog)
│   ├── blog_post_draft.md  # Blog post
│   ├── COMPILED_LANGUAGES.md  # Compiled language integration
│   ├── MODULAR_FORUM_ANNOUNCEMENT.md
│   └── MARIMO_ANNOUNCEMENT.md
└── README.md                # This file

Development

Using just (Recommended)

We provide a justfile with common tasks synced across uv and pixi:

# Show all available commands
just --list

# Install dependencies
just install

# Run tests
just test
just test-coverage

# Code quality
just format
just lint
just typecheck
just check              # Run all quality checks

# Notebooks
just notebook-decorator
just notebook-executor
just benchmark

# Development
just clean
just clean-mojo-cache
just cache-stats

# CI checks locally
just ci

Using uv

# Install in development mode
uv sync --extra dev

# Run tests
uv run pytest tests/

# Format code
uv run ruff format .

# Lint
uv run ruff check .

# Type check
uv run ty check

Using pixi

# Run tests
pixi run test

# Format code
pixi run format

# Lint
pixi run lint

# Type check
pixi run typecheck

# All quality checks
pixi run check

Why marimo?

marimo is a reactive Python notebook with several advantages:

  • Reactive execution: Change a slider, Mojo re-runs automatically
  • Pure Python files: Version control friendly (unlike Jupyter JSON)
  • Type-safe: Built-in UI elements with proper types
  • Reproducible: Dependency graph prevents hidden state bugs

Contributing

Contributions welcome! This is a living experiment, evolving based on real-world usage.

Areas for contribution:

  • Additional notebook examples
  • Performance profiling on diverse hardware
  • Integration with other notebook environments
  • Error handling improvements
  • Documentation enhancements

Related Projects

License

Apache License 2.0 - See LICENSE for details.

About

Created by Michael Booth at DataBooth.

Part of an ongoing exploration of Mojo for real-world data and AI workflows. See the full series at databooth.com.au/posts/mojo.

DataBooth helps medium-sized businesses leverage high-performance computing for data analytics and AI, now offering Mojo-powered services that deliver 10-100× faster solutions without vendor lock-in.

About

Interactive Mojo integration for Python notebooks - three patterns for running high-performance Mojo code from marimo

Topics

Resources

License

Stars

Watchers

Forks

Releases

No releases published

Packages

No packages published

Contributors 2

  •  
  •