Artifact API

mlx_lattice.artifact loads and saves legacy lattice model artifacts. The deployment entry point is mlx_lattice.artifact.load_lattice_model(), which reads a manifest and safetensors weight file, validates the graph, and returns an in-memory mlx_lattice.artifact.LatticeModel.

The public package root intentionally exposes only loading, saving, and the runtime model. JSON graph builders and registries are legacy submodule tools; future producer work should target the MLIR lattice dialect rather than growing the JSON graph authoring surface.

The artifact implementation has three layers:

Module artifact helpers under mlx_lattice.artifact.builder still build legacy manifests and weight dictionaries from approved sparse NN modules. They are kept for tests, local fixtures, and compatibility with the current JSON runner, but they are not the long-term cross-framework IR API.

The artifact runner honors manifest dtype_policy for floating dense arrays and sparse feature matrices while preserving coordinates, integer arrays, byte streams, and packed quantized payloads.

Artifact manifests also carry runtime compatibility metadata. The loader accepts artifacts targeted at mlx-lattice whose version specifier matches the installed native mlx-lattice package, and rejects incompatible runtime metadata names or version windows before dispatching any graph node.

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

Bases: object

Loaded lattice model artifact.

Parameters:
manifest: IRManifest
weights: dict[str, array]
model()[source]

Construct an executable in-memory MLX graph.

Return type:

LatticeModel

class mlx_lattice.artifact.LatticeModel(manifest, weights)[source]

Bases: object

Executable in-memory graph loaded from a lattice artifact.

The model does not execute generated Python. It validates IR nodes and dispatches each node through the public mlx_lattice.ops surface.

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

Load a lattice artifact directory.

The artifact directory must contain manifest.json and weights.safetensors. Use LatticeArtifact.model() or load_lattice_model() when an executable graph runner is needed.

Return type:

LatticeArtifact

Parameters:

path (str | Path)

mlx_lattice.artifact.load_lattice_model(path)[source]

Load a lattice artifact directory into an executable model.

The model-oriented name is the recommended deployment entry point.

Return type:

LatticeModel

Parameters:

path (str | Path)

mlx_lattice.artifact.save_lattice_artifact(path, manifest, weights)[source]

Write a lattice artifact directory.

This helper is intentionally strict and low level. It is useful for tests, fixtures, and future producers that already own a valid manifest.

Return type:

None

Parameters:
mlx_lattice.artifact.save_lattice_graph(path, builder, **kwargs)[source]

Build and save an explicitly built lattice graph.

Return type:

None

Parameters:
mlx_lattice.artifact.save_lattice_model(path, manifest, weights)[source]

Alias for save_lattice_artifact().

Return type:

None

Parameters:
mlx_lattice.artifact.save_lattice_module(path, module, **kwargs)[source]

Build and save a serializable sparse NN module.

Return type:

None

Parameters:

  • path (str | Path)

  • module (Module)

Artifact I/O

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

Bases: object

Loaded lattice model artifact.

Parameters:
manifest: IRManifest
weights: dict[str, array]
model()[source]

Construct an executable in-memory MLX graph.

Return type:

LatticeModel

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

Load a lattice artifact directory.

The artifact directory must contain manifest.json and weights.safetensors. Use LatticeArtifact.model() or load_lattice_model() when an executable graph runner is needed.

Return type:

LatticeArtifact

Parameters:

path (str | Path)

mlx_lattice.artifact.io.load_lattice_model(path)[source]

Load a lattice artifact directory into an executable model.

The model-oriented name is the recommended deployment entry point.

Return type:

LatticeModel

Parameters:

path (str | Path)

mlx_lattice.artifact.io.save_lattice_artifact(path, manifest, weights)[source]

Write a lattice artifact directory.

This helper is intentionally strict and low level. It is useful for tests, fixtures, and future producers that already own a valid manifest.

Return type:

None

Parameters:
mlx_lattice.artifact.io.save_lattice_model(path, manifest, weights)[source]

Alias for save_lattice_artifact().

Return type:

None

Parameters:
mlx_lattice.artifact.io.save_lattice_module(path, module, **kwargs)[source]

Build and save a serializable sparse NN module.

Return type:

None

Parameters:

  • path (str | Path)

  • module (Module)

mlx_lattice.artifact.io.save_lattice_graph(path, builder, **kwargs)[source]

Build and save an explicitly built lattice graph.

Return type:

None

Parameters:

Artifact model

class mlx_lattice.artifact.model.LatticeModel(manifest, weights)[source]

Bases: object

Executable in-memory graph loaded from a lattice artifact.

The model does not execute generated Python. It validates IR nodes and dispatches each node through the public mlx_lattice.ops surface.

Parameters:

Legacy internals

mlx_lattice.artifact.registry.lattice_op(op, *, function, inputs, outputs=None, parameters=None, optional_parameters=None, attributes=None, value_attributes=None, defaults=None, sequence_inputs=None, output='output', output_types=None, input_types=None, value_attribute_types=None, handler=None)[source]

Register a public operation with one compact annotation.

Return type:

Callable[[TypeVar(HandlerT, bound= Callable[[ExecutionContext, IRNode], dict[str, SparseTensor | array | KernelRelation | NeighborRelation | CoordinateSet | SparseAlignment | SparseQuantization | PointVoxelMap | CoordinateOrdering | SparseOccupancy | OccupancyExpansion | bytes | object]])], TypeVar(HandlerT, bound= Callable[[ExecutionContext, IRNode], dict[str, SparseTensor | array | KernelRelation | NeighborRelation | CoordinateSet | SparseAlignment | SparseQuantization | PointVoxelMap | CoordinateOrdering | SparseOccupancy | OccupancyExpansion | bytes | object]])]

Parameters:

  • op (str | IROpContract)

  • function (Callable[[...], GraphValue])

  • inputs (Mapping[str, str])

  • outputs (set[str] | None)

  • parameters (Mapping[str, str | ParameterBinding] | None)

  • optional_parameters (Mapping[str, str | ParameterBinding] | None)

  • attributes (Mapping[str, str] | None)

  • value_attributes (Mapping[str, str] | None)

  • defaults (Mapping[str, Any] | None)

  • sequence_inputs (set[str] | None)

  • output (str)

  • output_types (Mapping[str, Literal['any', 'sparse_tensor', 'dense_tensor', 'relation', 'coordinate_set', 'alignment', 'quantization', 'point_voxel_map', 'coordinate_ordering', 'sparse_occupancy', 'occupancy_expansion', 'bytes']] | None)

  • input_types (Mapping[str, Literal['any', 'sparse_tensor', 'dense_tensor', 'relation', 'coordinate_set', 'alignment', 'quantization', 'point_voxel_map', 'coordinate_ordering', 'sparse_occupancy', 'occupancy_expansion', 'bytes']] | None)

  • value_attribute_types (Mapping[str, Literal['any', 'sparse_tensor', 'dense_tensor', 'relation', 'coordinate_set', 'alignment', 'quantization', 'point_voxel_map', 'coordinate_ordering', 'sparse_occupancy', 'occupancy_expansion', 'bytes']] | None)

  • handler (GraphHandler | None)

mlx_lattice.artifact.registry.module_binding(module_type, *, op, parameters=(), attributes=<function <lambda>>)[source]

Register a serializable NN module producer.

Return type:

Callable[[type[TypeVar(ModuleT, bound= Module)]], type[TypeVar(ModuleT, bound= Module)]]

Parameters:
mlx_lattice.artifact.registry.iter_operation_specs()[source]

Return all operation contracts supported by the MLX artifact graph.

Return type:

tuple[IROpSpec, ...]

mlx_lattice.artifact.registry.operation_spec(name)[source]

Return the artifact operation spec for name.

Return type:

IROpSpec

Parameters:

name (str)

mlx_lattice.artifact.registry.operation_binding(name)[source]

Return the artifact operation binding for name.

Return type:

OperationBinding

Parameters:

name (str)

mlx_lattice.artifact.registry.module_artifact_binding(module)[source]

Return the artifact binding for module.

Return type:

ModuleBinding

Parameters:

module (Module)

mlx_lattice.artifact.registry.validate_node_against_artifact(node)[source]

Validate an IR node against the MLX artifact binding.

Return type:

None

Parameters:

node (IRNode)

type mlx_lattice.artifact.bindings.GraphValue = SparseTensor | array | KernelRelation | NeighborRelation | CoordinateSet | SparseAlignment | SparseQuantization | PointVoxelMap | CoordinateOrdering | SparseOccupancy | OccupancyExpansion | bytes | object
type mlx_lattice.artifact.bindings.GraphHandler = Callable[[ExecutionContext, IRNode], dict[str, GraphValue]]
class mlx_lattice.artifact.bindings.ParameterBinding(argument, kind=IRParameterKind.ARRAY)[source]

Bases: object

How an IR parameter name is converted into a runtime argument.

Parameters:
argument: str
kind: IRParameterKind
class mlx_lattice.artifact.bindings.ModuleParameterBinding(name, source)[source]

Bases: object

How an NN module field is stored as an IR parameter.

Parameters:
name: str
source: str
class mlx_lattice.artifact.bindings.ValueTypeBinding(value_type, runtime_type, fields=<factory>)[source]

Bases: object

Runtime type and public fields for one IR value type.

Parameters:

  • value_type (Literal['any', 'sparse_tensor', 'dense_tensor', 'relation', 'coordinate_set', 'alignment', 'quantization', 'point_voxel_map', 'coordinate_ordering', 'sparse_occupancy', 'occupancy_expansion', 'bytes'])

  • runtime_type (type | tuple[type, ...])

  • fields (Mapping[str, Literal['any', 'sparse_tensor', 'dense_tensor', 'relation', 'coordinate_set', 'alignment', 'quantization', 'point_voxel_map', 'coordinate_ordering', 'sparse_occupancy', 'occupancy_expansion', 'bytes']])

value_type: Literal['any', 'sparse_tensor', 'dense_tensor', 'relation', 'coordinate_set', 'alignment', 'quantization', 'point_voxel_map', 'coordinate_ordering', 'sparse_occupancy', 'occupancy_expansion', 'bytes']
runtime_type: type | tuple[type, ...]
fields: Mapping[str, Literal['any', 'sparse_tensor', 'dense_tensor', 'relation', 'coordinate_set', 'alignment', 'quantization', 'point_voxel_map', 'coordinate_ordering', 'sparse_occupancy', 'occupancy_expansion', 'bytes']]
class mlx_lattice.artifact.bindings.OperationBinding(spec, function, input_arguments, parameter_arguments=<factory>, attribute_arguments=<factory>, value_attribute_arguments=<factory>, defaults=<factory>, sequence_inputs=frozenset({}), output='output', handler=None)[source]

Bases: object

Runtime binding for one public lattice operation.

Parameters:
spec: IROpSpec
function: Callable[[...], GraphValue]
input_arguments: Mapping[str, str]
parameter_arguments: Mapping[str, ParameterBinding]
attribute_arguments: Mapping[str, str]
value_attribute_arguments: Mapping[str, str]
defaults: Mapping[str, Any]
sequence_inputs: frozenset[str]
output: str
handler: GraphHandler | None
run(context, node)[source]
Return type:

dict[str, SparseTensor | array | KernelRelation | NeighborRelation | CoordinateSet | SparseAlignment | SparseQuantization | PointVoxelMap | CoordinateOrdering | SparseOccupancy | OccupancyExpansion | bytes | object]

Parameters:
run_default(context, node)[source]
Return type:

dict[str, SparseTensor | array | KernelRelation | NeighborRelation | CoordinateSet | SparseAlignment | SparseQuantization | PointVoxelMap | CoordinateOrdering | SparseOccupancy | OccupancyExpansion | bytes | object]

Parameters:
arguments(context, node)[source]
Return type:

dict[str, Any]

Parameters:
class mlx_lattice.artifact.bindings.ModuleBinding(module_type, op, parameters=(), attributes=<function ModuleBinding.<lambda>>)[source]

Bases: object

Artifact binding for a serializable sparse NN module.

Parameters:
module_type: type[Module]
op: str
parameters: tuple[ModuleParameterBinding, ...]
attributes: Callable[[Module], dict[str, Any]]
class mlx_lattice.artifact.bindings.ExecutionContext(values, weights, batch_size=None)[source]

Bases: object

Mutable graph execution state.

Parameters:
values: dict[str, GraphValue]
weights: Mapping[str, array]
batch_size: int | None
value(name)[source]
Return type:

SparseTensor | array | KernelRelation | NeighborRelation | CoordinateSet | SparseAlignment | SparseQuantization | PointVoxelMap | CoordinateOrdering | SparseOccupancy | OccupancyExpansion | bytes | object

Parameters:

name (str)

input_value(value_ref, *, sequence=False)[source]
Return type:

SparseTensor | array | KernelRelation | NeighborRelation | CoordinateSet | SparseAlignment | SparseQuantization | PointVoxelMap | CoordinateOrdering | SparseOccupancy | OccupancyExpansion | bytes | object | tuple[SparseTensor | array | KernelRelation | NeighborRelation | CoordinateSet | SparseAlignment | SparseQuantization | PointVoxelMap | CoordinateOrdering | SparseOccupancy | OccupancyExpansion | bytes | object, ...]

Parameters:
sparse(name)[source]
Return type:

SparseTensor

Parameters:

name (str)

parameter(name, kind=IRParameterKind.ARRAY)[source]
Return type:

Any

Parameters:
array(name)[source]
Return type:

array

Parameters:

name (str)

optional_array(name)[source]
Return type:

array | None

Parameters:

name (str | None)

quantized_weight(prefix)[source]
Return type:

QuantizedWeight

Parameters:

prefix (str)

mlx_lattice.artifact.bindings.infer_batch_size(values)[source]

Infer batch size from sparse graph inputs when metadata exists.

Return type:

int | None

Parameters:

values (Mapping[str, GraphValue])

mlx_lattice.artifact.bindings.apply_dtype_policy(value, policy, *, inference=True)[source]

Apply an artifact dtype policy to graph-carried tensor values.

Return type:

SparseTensor | array | KernelRelation | NeighborRelation | CoordinateSet | SparseAlignment | SparseQuantization | PointVoxelMap | CoordinateOrdering | SparseOccupancy | OccupancyExpansion | bytes | object

Parameters:
mlx_lattice.artifact.bindings.validate_value_type(name, expected, value)[source]

Validate a runtime value against an IR value type.

Return type:

None

Parameters:

  • name (str)

  • expected (Literal['any', 'sparse_tensor', 'dense_tensor', 'relation', 'coordinate_set', 'alignment', 'quantization', 'point_voxel_map', 'coordinate_ordering', 'sparse_occupancy', 'occupancy_expansion', 'bytes'])

  • value (GraphValue)

mlx_lattice.artifact.bindings.layout_id(value)[source]

Encode a quantized weight layout for artifact metadata.

Return type:

int

Parameters:

value (str)

mlx_lattice.artifact.bindings.value_type_binding(value_type)[source]

Return the artifact binding for an IR value type.

Return type:

ValueTypeBinding | None

Parameters:

value_type (Literal['any', 'sparse_tensor', 'dense_tensor', 'relation', 'coordinate_set', 'alignment', 'quantization', 'point_voxel_map', 'coordinate_ordering', 'sparse_occupancy', 'occupancy_expansion', 'bytes'])

mlx_lattice.artifact.bindings.value_type_fields(value_type)[source]

Return public structural fields exposed for an IR value type.

Return type:

Mapping[str, Literal['any', 'sparse_tensor', 'dense_tensor', 'relation', 'coordinate_set', 'alignment', 'quantization', 'point_voxel_map', 'coordinate_ordering', 'sparse_occupancy', 'occupancy_expansion', 'bytes']]

Parameters:

value_type (Literal['any', 'sparse_tensor', 'dense_tensor', 'relation', 'coordinate_set', 'alignment', 'quantization', 'point_voxel_map', 'coordinate_ordering', 'sparse_occupancy', 'occupancy_expansion', 'bytes'])

mlx_lattice.artifact.bindings.iter_value_type_bindings()[source]

Return artifact bindings for all known IR value types.

Return type:

tuple[ValueTypeBinding, ...]

class mlx_lattice.artifact.builder.LatticeArtifactData(manifest, weights)[source]

Bases: object

Manifest and weight tensors for an in-memory lattice artifact.

Parameters:
manifest: IRManifest
weights: dict[str, array]
class mlx_lattice.artifact.builder.GraphOutput(value, value_type=None, name=None)[source]

Bases: object

Public graph output mapping for an explicit artifact graph.

Parameters:

  • value (str)

  • value_type (Literal['any', 'sparse_tensor', 'dense_tensor', 'relation', 'coordinate_set', 'alignment', 'quantization', 'point_voxel_map', 'coordinate_ordering', 'sparse_occupancy', 'occupancy_expansion', 'bytes'] | None)

  • name (str | None)

value: str
value_type: Literal['any', 'sparse_tensor', 'dense_tensor', 'relation', 'coordinate_set', 'alignment', 'quantization', 'point_voxel_map', 'coordinate_ordering', 'sparse_occupancy', 'occupancy_expansion', 'bytes'] | None
name: str | None
class mlx_lattice.artifact.builder.LatticeGraphBuildable(*args, **kwargs)[source]

Bases: Protocol

Protocol for modules that append their own artifact graph nodes.

build_lattice_graph(builder, input_name)[source]

Append nodes to builder and return the output value name.

Return type:

str

Parameters:
class mlx_lattice.artifact.builder.LatticeGraphBuilder(input_name='input', *, inputs=None)[source]

Bases: object

Builder for explicit lattice model manifests.

The builder is intentionally small: it owns graph nodes and weight tensors, validates op/module bindings through the shared artifact registry, and returns named values that can be wired into later nodes. It is the escape hatch for custom modules and DAGs without requiring Python tracing.

Parameters:

  • input_name (str)

  • inputs (Mapping[str, IRValueType] | None)

value_types: dict[str, Literal['any', 'sparse_tensor', 'dense_tensor', 'relation', 'coordinate_set', 'alignment', 'quantization', 'point_voxel_map', 'coordinate_ordering', 'sparse_occupancy', 'occupancy_expansion', 'bytes']]
nodes: list[IRNode]
weights: dict[str, array]
add_op(name, op, *, inputs, output=None, outputs=None, parameters=None, attributes=None)[source]

Append an operation node and return its primary output value.

Return type:

str

Parameters:
call(op, /, name=None, *, output=None, outputs=None, parameters=None, **arguments)[source]

Append an operation by using its registered public argument names.

Graph-value arguments are passed as value names. JSON-compatible constants become manifest attributes. Tensor parameters are passed either as artifact key strings or as arrays/packed QuantizedWeight objects, which are stored automatically in weights.safetensors. The parameters mapping remains available when the caller already owns stable artifact keys. This is the concise explicit-graph API; add_op() remains the lower-level escape hatch when a manifest needs exact port dictionaries.

Return type:

str

Parameters:
add_module(name, module, input_value, *, output=None)[source]

Append a registered sparse NN module node.

Return type:

str

Parameters:

  • name (str)

  • module (Module)

  • input_value (str)

  • output (str | None)

output(value, *, name=None, value_type=None)[source]

Describe a public graph output, inferring its type by default.

Return type:

GraphOutput

Parameters:

  • value (str)

  • name (str | None)

  • value_type (Literal['any', 'sparse_tensor', 'dense_tensor', 'relation', 'coordinate_set', 'alignment', 'quantization', 'point_voxel_map', 'coordinate_ordering', 'sparse_occupancy', 'occupancy_expansion', 'bytes'] | None)

field(value, field, *, name=None)[source]

Project a supported structural field from a graph value.

Return type:

str

Parameters:
add_parameter(name, value)[source]

Store a dense tensor parameter and return its artifact key.

Return type:

str

Parameters:

  • name (str)

  • value (array)

add_quantized_parameter(name, value)[source]

Store a packed quantized parameter and return its artifact prefix.

Return type:

str

Parameters:
module_parameters(node_id, module, parameters)[source]

Artifact registered module parameters into artifact weights.

Return type:

dict[str, str]

Parameters:
unique_name(name)[source]

Return a graph-unique sanitized node name.

Return type:

str

Parameters:

name (str)

unique_weight_name(name)[source]

Return an artifact-unique sanitized weight name.

Return type:

str

Parameters:

name (str)

manifest(*, outputs=None, output_name='output', output_value=None, input_type=None, output_type=None, producer=None)[source]

Build an immutable manifest from the accumulated graph.

Return type:

IRManifest

Parameters:

  • outputs (Mapping[str, Literal['any', 'sparse_tensor', 'dense_tensor', 'relation', 'coordinate_set', 'alignment', 'quantization', 'point_voxel_map', 'coordinate_ordering', 'sparse_occupancy', 'occupancy_expansion', 'bytes'] | ~mlx_lattice.artifact.builder.GraphOutput | None] | ~collections.abc.Sequence[str] | None)

  • output_name (str)

  • output_value (str | None)

  • input_type (Literal['any', 'sparse_tensor', 'dense_tensor', 'relation', 'coordinate_set', 'alignment', 'quantization', 'point_voxel_map', 'coordinate_ordering', 'sparse_occupancy', 'occupancy_expansion', 'bytes'] | None)

  • output_type (Literal['any', 'sparse_tensor', 'dense_tensor', 'relation', 'coordinate_set', 'alignment', 'quantization', 'point_voxel_map', 'coordinate_ordering', 'sparse_occupancy', 'occupancy_expansion', 'bytes'] | None)

  • producer (Mapping[str, str] | None)

mlx_lattice.artifact.builder.build_lattice_graph_artifact(builder, *, outputs, producer=None)[source]

Build an artifact from an explicit lattice graph.

Return type:

LatticeArtifactData

Parameters:

  • builder (LatticeGraphBuilder)

  • outputs (Mapping[str, Literal['any', 'sparse_tensor', 'dense_tensor', 'relation', 'coordinate_set', 'alignment', 'quantization', 'point_voxel_map', 'coordinate_ordering', 'sparse_occupancy', 'occupancy_expansion', 'bytes'] | ~mlx_lattice.artifact.builder.GraphOutput | None] | ~collections.abc.Sequence[str])

  • producer (Mapping[str, str] | None)

mlx_lattice.artifact.builder.build_lattice_module_artifact(module, *, input_name='input', output_name='output', input_type='sparse_tensor', output_type=None, producer=None)[source]

Build a sparse NN module artifact graph.

Built-in lattice modules and sequential containers build structurally. Custom modules can implement build_lattice_graph(builder, input_name) to emit arbitrary DAGs with the same builder used internally.

Return type:

LatticeArtifactData

Parameters:

  • module (Module)

  • input_name (str)

  • output_name (str)

  • input_type (Literal['any', 'sparse_tensor', 'dense_tensor', 'relation', 'coordinate_set', 'alignment', 'quantization', 'point_voxel_map', 'coordinate_ordering', 'sparse_occupancy', 'occupancy_expansion', 'bytes'])

  • output_type (Literal['any', 'sparse_tensor', 'dense_tensor', 'relation', 'coordinate_set', 'alignment', 'quantization', 'point_voxel_map', 'coordinate_ordering', 'sparse_occupancy', 'occupancy_expansion', 'bytes'] | None)

  • producer (Mapping[str, str] | None)