mlc-python

MLC-Python

• :inbox_tray: Installation
• :key: Key Features
  • :building_construction: MLC Dataclass
  • :dart: Structure-Aware Tooling
  • :snake: Text Formats in Python AST
  • :zap: Zero-Copy Interoperability with C++ Plugins
• :fuelpump: Development
  • :gear: Editable Build
  • :ferris_wheel: Create Wheels

MLC is a Python-first toolkit that makes it more ergonomic to build AI compilers, runtimes, and compound AI systems with Pythonic dataclass, rich tooling infra and zero-copy interoperability with C++ plugins.

:inbox_tray: Installation

pip install -U mlc-python

:key: Key Features

:building_construction: MLC Dataclass

MLC dataclass is similar to Python’s native dataclass:

import mlc.dataclasses as mlcd

@mlcd.py_class("demo.MyClass")
class MyClass(mlcd.PyClass):
  a: int
  b: str
  c: float | None

instance = MyClass(12, "test", c=None)

Type safety. MLC dataclass checks type strictly in Cython and C++.

>>> instance.c = 10; print(instance)
demo.MyClass(a=12, b='test', c=10.0)

>>> instance.c = "wrong type"
TypeError: must be real number, not str

>>> instance.non_exist = 1
AttributeError: 'MyClass' object has no attribute 'non_exist' and no __dict__ for setting new attributes

Serialization. MLC dataclasses are picklable and JSON-serializable.

>>> MyClass.from_json(instance.json())
demo.MyClass(a=12, b='test', c=None)

>>> import pickle; pickle.loads(pickle.dumps(instance))
demo.MyClass(a=12, b='test', c=None)

:dart: Structure-Aware Tooling

An extra structure field are used to specify a dataclass's structure, indicating def site and scoping in an IR.

import mlc.dataclasses as mlcd

@mlcd.py_class
class Expr(mlcd.PyClass):
  def __add__(self, other):
    return Add(a=self, b=other)

@mlcd.py_class(structure="nobind")
class Add(Expr):
  a: Expr
  b: Expr

@mlcd.py_class(structure="var")
class Var(Expr):
  name: str = mlcd.field(structure=None) # excludes `name` from defined structure

@mlcd.py_class(structure="bind")
class Let(Expr):
  rhs: Expr
  lhs: Var = mlcd.field(structure="bind") # `Let.lhs` is the def-site
  body: Expr

Structural equality. Member method eq_s compares the structural equality (alpha equivalence) of two IRs represented by MLC's structured dataclass.

"""
L1: let z = x + y; z
L2: let x = y + z; x
L3: let z = x + x; z
"""
>>> x, y, z = Var("x"), Var("y"), Var("z")
>>> L1 = Let(rhs=x + y, lhs=z, body=z)
>>> L2 = Let(rhs=y + z, lhs=x, body=x)
>>> L3 = Let(rhs=x + x, lhs=z, body=z)
>>> L1.eq_s(L2)
True
>>> L1.eq_s(L3, assert_mode=True)
ValueError: Structural equality check failed at {root}.rhs.b: Inconsistent binding. RHS has been bound to a different node while LHS is not bound

Structural hashing. The structure of MLC dataclasses can be hashed via hash_s, which guarantees if two dataclasses are alpha-equivalent, they will share the same structural hash:

>>> L1_hash, L2_hash, L3_hash = L1.hash_s(), L2.hash_s(), L3.hash_s()
>>> assert L1_hash == L2_hash
>>> assert L1_hash != L3_hash

:snake: Text Formats in Python AST

TBD

:zap: Zero-Copy Interoperability with C++ Plugins

TBD

:fuelpump: Development

:gear: Editable Build

pip install --verbose --editable ".[dev]"
pre-commit install

:ferris_wheel: Create Wheels

This project uses cibuildwheel to build cross-platform wheels. See .github/workflows/wheels.ym for more details.

export CIBW_BUILD_VERBOSITY=3
export CIBW_BUILD="cp3*-manylinux_x86_64"
python -m pip install pipx
pipx run cibuildwheel==2.20.0 --output-dir wheelhouse