Artifact API

mlx_lattice.artifact handles the MLX-side consumer half of the portable model-artifact boundary. A bundle contains textual graph.mlir and weights.safetensors. The bundle itself is exchange media; execution starts only after load_lattice_program or compile_lattice_artifact validates and lowers the graph.

The producer half belongs to training frameworks. For example, the Torch/CUDA side should expose model-to-artifact APIs such as save_lattice_model_artifact. Those APIs lower a framework model into graph.mlir plus weights. They should not be confused with MLX raw artifact IO, which only saves, loads, validates, and compiles an already-formed bundle.

MLIR is the only artifact graph contract. Published macOS wheels include native MLIR execution support. A source build may still lack that support when it was compiled without MLX_LATTICE_ENABLE_MLIR. In that case bundle IO remains available, validation can still run through lattice-opt when provided, and native_artifact_execution_available() returns False. compile_lattice_artifact and load_lattice_program then fail with a direct capability error instead of selecting another graph runtime.

Graph validation uses real MLIR infrastructure. MLIR-enabled native builds use an in-process parser/verifier registered with the lattice dialect. Lightweight builds can validate through the repository lattice-opt tool. The Python layer does not parse lattice MLIR text itself.

Artifact validation also enforces the module-level ABI contract: lattice.ir_version must be 2 and lattice.weight_file must be "weights.safetensors". Unsupported versions or payload names fail before the graph reaches runtime lowering. The module must also contain exactly one func.func entry named forward and that function must return at least one value. The entry body is intentionally simple: one block containing only lattice dialect operations and func.return. This matches the v2 MLX importer contract and prevents silently ignored MLIR operations.

Executable lowering is available when the MLIR-enabled native extension is built. The native extension parses and verifies the module, returns a structured runtime plan, and Python lowers that plan through existing mlx_lattice.ops routes.

Runtime plan inputs carry importer-generated value names plus MLIR-derived metadata: type text and ABI role. SparseTensor shorthand binding uses the roles sparse_coords, sparse_features, and sparse_active on the first three entry arguments. It does not depend on textual MLIR SSA argument names, which are not a stable API contract.

Programs can be invoked with one SparseTensor shorthand argument, positional ABI tensors, or keyword ABI tensors:

program(x)
program(x.coords, x.feats, x.active_rows)
program(coords=x.coords, features=x.feats, active=x.active_rows)

These forms are intentionally strict. Sparse shorthand cannot be mixed with keyword inputs; missing, duplicate, unexpected, or non-MLX input values fail before any operation lowering runs. This makes producer/consumer drift visible at the artifact boundary instead of inside a Metal kernel or sparse operator.

The raw native payload is frozen into RuntimePlan before execution. LatticeProgram consumes that typed plan object, not arbitrary dictionaries, so malformed importer payloads fail before any MLX operation lowering runs. Plan freezing is schema-aware and ABI-aware: the plan must carry ir_version and weight_file values matching the verified module-level contract, operation names must resolve through lattice-contract, SSA operand/result counts must match the contract op definition, required attributes must be present, enum attributes and structured packing metadata must use valid v2 values, triple/numeric/string attributes must have the expected payload shape, runtime value labels must be defined before use, and returns must reference defined values. Frozen plans also deep-freeze nested attribute payloads such as packing maps and stride triples so runtime lowerings cannot accidentally mutate importer state.

The MLX importer lowers the sparse ABI operations, symbolic dense and packed affine int4/int8 weights, forward/submanifold/target/transpose/generative sparse convolution, local/global sparse pooling, point/voxel conversion, feature projection, feature activation, sparse feature replacement, sparse feature normalization, and sparse binary algebra. Unsupported dialect operations should fail during lowering instead of being interpreted from ad-hoc JSON or Python-side string switches.

Runtime lowerings are schema-bound. A lowering names parameters after lattice operands and attributes, and artifact_lowering binds those parameters from the dialect schema. Operation modules do not manually inspect PlanOperation.operands or PlanOperation.attrs. Extra metadata is only needed for semantic cases that the raw MLIR type cannot express by itself, such as resolving a packed convolution weight with the input channel count. The intended lowering shape is a small semantic bridge:

@artifact_lowering(op=conv3d, weights={"weight": conv_weight(input="input")})
def conv3d_from_artifact(
    input: SparseTensor,
    weight: mx.array | QuantizedWeight,
    bias: mx.array | None = None,
    *,
    kernel_size: Triple,
    stride: Triple,
    padding: Triple,
    dilation: Triple,
) -> SparseTensor:
    return conv3d(
        input,
        weight,
        bias,
        kernel_size=kernel_size,
        stride=stride,
        padding=padding,
        dilation=dilation,
    )

Dense weights are read from weights.safetensors using the exact storage_key declared by lattice.weight. Packed weights use <storage_key>.weight, <storage_key>.scales, and <storage_key>.biases. Bias parameters are regular dense lattice.weight values with family bias and layout bias_c; they are passed as optional operands to convolution-family and linear operations. Missing dense weights, incomplete packed triplets, or packed scale dtype drift are runtime contract errors.

Feature-only dialect operations operate on dense feature tensors. This includes lattice.linear, lattice.activation, lattice.batch_norm, lattice.layer_norm, and lattice.rms_norm. Sparse support preservation must be explicit in MLIR through lattice.sparse.decompose followed by lattice.sparse.with_features.

Bundle IO

class mlx_lattice.artifact.LatticeArtifact(graph, weights)[source]

Bases: object

MLIR-first lattice artifact bundle.

The artifact is intentionally just the exchange media at this layer: textual lattice MLIR plus safetensors weights. Executable MLX import is a separate lowering step built on top of the MLIR contract.

Parameters:
graph: str
weights: dict[str, array]
class mlx_lattice.artifact.LatticeMLIRStatus(valid, diagnostics='')[source]

Bases: object

Native MLIR validation result for a lattice graph.

Parameters:
valid: bool
diagnostics: str
class mlx_lattice.artifact.LatticeProgram(plan, weights)[source]

Bases: object

Executable MLX lowering of a verified lattice MLIR artifact.

The plan is produced by the MLIR-enabled native extension after parsing and verifying the graph with the registered lattice dialect. Python only lowers the structured plan to existing mlx_lattice.ops calls.

Parameters:
plan: RuntimePlan
weights: Mapping[str, array]
classmethod from_artifact(artifact)[source]

Compile a loaded artifact through the native MLIR importer.

Return type:

LatticeProgram

Parameters:

artifact (LatticeArtifact)

call_with_trusted_sparse_inputs(*args, **kwargs)[source]

Execute with stride-qualified sparse coordinates.

This inference ABI skips host validation of sparse coordinate stride divisibility. Callers must establish that contract before execution.

Return type:

SparseTensor | array | PackedWeightPayload | tuple[SparseTensor | array | PackedWeightPayload, ...]

Parameters:
class mlx_lattice.artifact.PlanArgument(name, abi_name, type, role)[source]

Bases: object

One verified artifact entry argument.

name is an importer-generated runtime value label. type is the canonical MLIR type text. role is MLIR-derived ABI metadata used for ergonomic bindings such as SparseTensor shorthand.

Parameters:
name: str
abi_name: str
type: str
role: str
class mlx_lattice.artifact.PlanOperation(name, definition, operands, results, attrs, operand_types, result_types)[source]

Bases: object

One executable lattice operation in native importer order.

Parameters:
name: str
definition: OpDef
operands: tuple[str, ...]
results: tuple[str, ...]
attrs: Mapping[str, Any]
operand_types: tuple[str, ...]
result_types: tuple[str, ...]
class mlx_lattice.artifact.PlanOutput(name, abi_name, type, role)[source]

Bases: object

One verified artifact entry output ABI item.

Parameters:
name: str
abi_name: str
type: str
role: str
class mlx_lattice.artifact.RuntimePlan(ir_version, schema_digest, weight_file, name, args, ops, returns, outputs)[source]

Bases: object

Typed Python view of a native-verified lattice MLIR execution plan.

Parameters:
ir_version: int
schema_digest: str
weight_file: str
name: str
args: tuple[PlanArgument, ...]
ops: tuple[PlanOperation, ...]
returns: tuple[str, ...]
outputs: tuple[PlanOutput, ...]
classmethod from_native(raw)[source]

Validate and freeze a native lattice_mlir_plan payload.

Return type:

RuntimePlan

Parameters:

raw (Mapping[str, Any])

mlx_lattice.artifact.compile_lattice_artifact(artifact)[source]

Compile a loaded artifact into an executable MLX program.

Return type:

LatticeProgram

Parameters:

artifact (LatticeArtifact)

mlx_lattice.artifact.lattice_artifact_status(artifact, *, lattice_opt=None)[source]

Return MLIR validation status for a lattice artifact.

Return type:

LatticeMLIRStatus

Parameters:
mlx_lattice.artifact.lattice_graph_operation_names(graph)[source]

Return lattice operation names from a native-verified graph.

This requires an MLIR-enabled native extension. It is the first typed importer bridge and deliberately does not fall back to textual parsing.

Return type:

tuple[str, ...]

Parameters:

graph (str)

mlx_lattice.artifact.lattice_graph_status(graph, *, lattice_opt=None)[source]

Return MLIR parse/verify status for a lattice graph.

Return type:

LatticeMLIRStatus

Parameters:
mlx_lattice.artifact.load_lattice_artifact(path)[source]

Load a lattice MLIR artifact directory.

Return type:

LatticeArtifact

Parameters:

path (str | Path)

mlx_lattice.artifact.load_lattice_program(path)[source]

Load an artifact directory and compile it into an MLX program.

Return type:

LatticeProgram

Parameters:

path (str | Path)

mlx_lattice.artifact.native_artifact_execution_available()[source]

Return whether this install can execute MLIR artifacts natively.

Return type:

bool

mlx_lattice.artifact.save_lattice_artifact(path, graph, weights=None)[source]

Write a lattice MLIR artifact directory.

Return type:

None

Parameters:
mlx_lattice.artifact.validate_lattice_artifact(artifact, *, lattice_opt=None)[source]

Parse and verify the graph contained in a lattice artifact.

Return type:

None

Parameters:
mlx_lattice.artifact.validate_lattice_graph(graph, *, lattice_opt=None)[source]

Parse and verify lattice MLIR with native MLIR infrastructure.

MLIR-enabled native builds use the in-process parser/verifier binding. Lightweight builds can still validate through lattice-opt by passing a tool path, setting MLX_LATTICE_LATTICE_OPT, or building the repository clangd-mlir preset.

Return type:

None

Parameters:
class mlx_lattice.artifact.io.LatticeArtifact(graph, weights)[source]

Bases: object

MLIR-first lattice artifact bundle.

The artifact is intentionally just the exchange media at this layer: textual lattice MLIR plus safetensors weights. Executable MLX import is a separate lowering step built on top of the MLIR contract.

Parameters:
graph: str
weights: dict[str, array]
mlx_lattice.artifact.io.save_lattice_artifact(path, graph, weights=None)[source]

Write a lattice MLIR artifact directory.

Return type:

None

Parameters:
mlx_lattice.artifact.io.load_lattice_artifact(path)[source]

Load a lattice MLIR artifact directory.

Return type:

LatticeArtifact

Parameters:

path (str | Path)

Executable import

class mlx_lattice.artifact.runtime.LatticeProgram(plan, weights)[source]

Bases: object

Executable MLX lowering of a verified lattice MLIR artifact.

The plan is produced by the MLIR-enabled native extension after parsing and verifying the graph with the registered lattice dialect. Python only lowers the structured plan to existing mlx_lattice.ops calls.

Parameters:
plan: RuntimePlan
weights: Mapping[str, array]
classmethod from_artifact(artifact)[source]

Compile a loaded artifact through the native MLIR importer.

Return type:

LatticeProgram

Parameters:

artifact (LatticeArtifact)

call_with_trusted_sparse_inputs(*args, **kwargs)[source]

Execute with stride-qualified sparse coordinates.

This inference ABI skips host validation of sparse coordinate stride divisibility. Callers must establish that contract before execution.

Return type:

SparseTensor | array | PackedWeightPayload | tuple[SparseTensor | array | PackedWeightPayload, ...]

Parameters:
mlx_lattice.artifact.runtime.compile_lattice_artifact(artifact)[source]

Compile a loaded artifact into an executable MLX program.

Return type:

LatticeProgram

Parameters:

artifact (LatticeArtifact)

mlx_lattice.artifact.runtime.load_lattice_program(path)[source]

Load an artifact directory and compile it into an MLX program.

Return type:

LatticeProgram

Parameters:

path (str | Path)

mlx_lattice.artifact.runtime.native_artifact_execution_available()[source]

Return whether this install can execute MLIR artifacts natively.

Return type:

bool

mlx_lattice.artifact.runtime.weight_from_artifact(program, storage_key, layout, packing)[source]
Return type:

array | PackedWeightPayload

Parameters:
mlx_lattice.artifact.runtime.sparse_make_from_artifact(program, coords, features, active, *, stride, coord_order)[source]
Return type:

SparseTensor

Parameters:

  • program (_ArtifactExecutionContext)

  • coords (array)

  • features (array)

  • active (array)

  • stride (tuple[int, int, int])

  • coord_order (str)

mlx_lattice.artifact.runtime.sparse_decompose_from_artifact(input)[source]
Return type:

tuple[array, array, array]

Parameters:

input (SparseTensor)

Runtime plan

class mlx_lattice.artifact.plan.PlanArgument(name, abi_name, type, role)[source]

Bases: object

One verified artifact entry argument.

name is an importer-generated runtime value label. type is the canonical MLIR type text. role is MLIR-derived ABI metadata used for ergonomic bindings such as SparseTensor shorthand.

Parameters:
name: str
abi_name: str
type: str
role: str
class mlx_lattice.artifact.plan.PlanOutput(name, abi_name, type, role)[source]

Bases: object

One verified artifact entry output ABI item.

Parameters:
name: str
abi_name: str
type: str
role: str
class mlx_lattice.artifact.plan.PlanOperation(name, definition, operands, results, attrs, operand_types, result_types)[source]

Bases: object

One executable lattice operation in native importer order.

Parameters:
name: str
definition: OpDef
operands: tuple[str, ...]
results: tuple[str, ...]
attrs: Mapping[str, Any]
operand_types: tuple[str, ...]
result_types: tuple[str, ...]
class mlx_lattice.artifact.plan.RuntimePlan(ir_version, schema_digest, weight_file, name, args, ops, returns, outputs)[source]

Bases: object

Typed Python view of a native-verified lattice MLIR execution plan.

Parameters:
ir_version: int
schema_digest: str
weight_file: str
name: str
args: tuple[PlanArgument, ...]
ops: tuple[PlanOperation, ...]
returns: tuple[str, ...]
outputs: tuple[PlanOutput, ...]
classmethod from_native(raw)[source]

Validate and freeze a native lattice_mlir_plan payload.

Return type:

RuntimePlan

Parameters:

raw (Mapping[str, Any])

MLIR validation

class mlx_lattice.artifact.validation.LatticeMLIRStatus(valid, diagnostics='')[source]

Bases: object

Native MLIR validation result for a lattice graph.

Parameters:
valid: bool
diagnostics: str
mlx_lattice.artifact.validation.validate_lattice_graph(graph, *, lattice_opt=None)[source]

Parse and verify lattice MLIR with native MLIR infrastructure.

MLIR-enabled native builds use the in-process parser/verifier binding. Lightweight builds can still validate through lattice-opt by passing a tool path, setting MLX_LATTICE_LATTICE_OPT, or building the repository clangd-mlir preset.

Return type:

None

Parameters:
mlx_lattice.artifact.validation.lattice_graph_status(graph, *, lattice_opt=None)[source]

Return MLIR parse/verify status for a lattice graph.

Return type:

LatticeMLIRStatus

Parameters:
mlx_lattice.artifact.validation.lattice_graph_operation_names(graph)[source]

Return lattice operation names from a native-verified graph.

This requires an MLIR-enabled native extension. It is the first typed importer bridge and deliberately does not fall back to textual parsing.

Return type:

tuple[str, ...]

Parameters:

graph (str)

mlx_lattice.artifact.validation.validate_lattice_artifact(artifact, *, lattice_opt=None)[source]

Parse and verify the graph contained in a lattice artifact.

Return type:

None

Parameters:
mlx_lattice.artifact.validation.lattice_artifact_status(artifact, *, lattice_opt=None)[source]

Return MLIR validation status for a lattice artifact.

Return type:

LatticeMLIRStatus

Parameters: