lattice-contract 0.1.0__tar.gz

lattice_contract-0.1.0/.gitignore

@@ -0,0 +1,27 @@
+# uv
+.venv
+.python-version
+.ruff_cache
+.pytest_cache
+__pycache__
+*.py[cod]
+dist
+build
+benchmarks/results/
+**/_ext*.so
+**/_ext*.pyd
+**/_ext*.dll
+**/_ext*.dylib
+**/_ext.pyi
+
+# ly's choice
+references
+
+# preferences
+/.vscode
+/.zed
+/.vercel
+.env*
+
+# os-related
+.DS_Store

lattice_contract-0.1.0/PKG-INFO

@@ -0,0 +1,21 @@
+Metadata-Version: 2.4
+Name: lattice-contract
+Version: 0.1.0
+Summary: Backend-neutral sparse lattice artifact contract
+Author-email: "Z.Y. Lin" <me@iki.moe>
+License-Expression: MIT
+Keywords: ir,lattice,model-artifact,sparse
+Requires-Python: >=3.12
+Description-Content-Type: text/markdown
+
+# lattice-contract
+
+`lattice-contract` is the backend-neutral artifact contract package used by
+`mlx-lattice`.
+
+It contains the sparse model manifest dataclasses, validation helpers, value
+type names, and operation-contract annotations. It intentionally does not
+import MLX, Torch, native extensions, or backend runtime objects.
+
+This package exists so training-side and deployment-side tools can share the
+same sparse graph contract while keeping their execution stacks independent.

lattice_contract-0.1.0/README.md

@@ -0,0 +1,11 @@
+# lattice-contract
+
+`lattice-contract` is the backend-neutral artifact contract package used by
+`mlx-lattice`.
+
+It contains the sparse model manifest dataclasses, validation helpers, value
+type names, and operation-contract annotations. It intentionally does not
+import MLX, Torch, native extensions, or backend runtime objects.
+
+This package exists so training-side and deployment-side tools can share the
+same sparse graph contract while keeping their execution stacks independent.

lattice_contract-0.1.0/lattice_contract/__init__.py

@@ -0,0 +1,62 @@
+from __future__ import annotations
+
+from lattice_contract.manifest import (
+    CURRENT_SCHEMA_VERSION,
+    DTypePolicy,
+    IRInputRef,
+    IRManifest,
+    IRNode,
+    IRParameter,
+    IRSparseSupport,
+    IRTensorSpec,
+    IRValueType,
+    Triple,
+    ir_value_type,
+    is_ir_value_type,
+    load_manifest,
+    manifest_from_dict,
+    manifest_to_dict,
+    triple,
+)
+from lattice_contract.ops import (
+    IROpArtifactHints,
+    IROpSpec,
+    IRParameterKind,
+    ir_op_spec,
+    iter_op_specs,
+    lattice_op_hints,
+    op_artifact_hints,
+    op_spec,
+    validate_node_against_spec,
+)
+
+__version__ = '0.1.0'
+
+__all__ = [
+    'CURRENT_SCHEMA_VERSION',
+    'DTypePolicy',
+    'IRInputRef',
+    'IRManifest',
+    'IRNode',
+    'IROpArtifactHints',
+    'IROpSpec',
+    'IRParameter',
+    'IRParameterKind',
+    'IRSparseSupport',
+    'IRTensorSpec',
+    'IRValueType',
+    'Triple',
+    '__version__',
+    'ir_op_spec',
+    'ir_value_type',
+    'is_ir_value_type',
+    'iter_op_specs',
+    'lattice_op_hints',
+    'load_manifest',
+    'manifest_from_dict',
+    'manifest_to_dict',
+    'op_artifact_hints',
+    'op_spec',
+    'triple',
+    'validate_node_against_spec',
+]

lattice_contract-0.1.0/lattice_contract/manifest.py

@@ -0,0 +1,456 @@
+from __future__ import annotations
+
+import json
+from collections.abc import Mapping, Sequence
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any, Literal, cast
+
+CURRENT_SCHEMA_VERSION = '0.1'
+
+DTypePolicy = Literal['preserve', 'fp32', 'fp16', 'fp16_inference']
+IRValueType = Literal[
+    'any',
+    'sparse_tensor',
+    'dense_tensor',
+    'relation',
+    'coordinate_set',
+    'alignment',
+    'quantization',
+    'point_voxel_map',
+    'coordinate_ordering',
+    'sparse_occupancy',
+    'occupancy_expansion',
+    'bytes',
+]
+IRParameter = str
+type IRInputRef = str | tuple[str, ...]
+Triple = tuple[int, int, int]
+
+_IR_VALUE_TYPES = frozenset(
+    (
+        'any',
+        'sparse_tensor',
+        'dense_tensor',
+        'relation',
+        'coordinate_set',
+        'alignment',
+        'quantization',
+        'point_voxel_map',
+        'coordinate_ordering',
+        'sparse_occupancy',
+        'occupancy_expansion',
+        'bytes',
+    )
+)
+_COORDINATE_ORDER = ('batch', 'x', 'y', 'z')
+_FEATURE_LAYOUT = ('row', 'channel')
+_WEIGHT_LAYOUT = 'mlx-lattice'
+
+
+def triple(value: int | Sequence[int], *, name: str) -> Triple:
+    """Normalize an integer or 3-sequence into a spatial integer triple."""
+
+    if isinstance(value, int):
+        return (value, value, value)
+    if len(value) != 3:
+        raise ValueError(f'{name} must be an int or a sequence of 3 ints.')
+    out = tuple(int(item) for item in value)
+    if len(out) != 3:
+        raise ValueError(f'{name} must be an int or a sequence of 3 ints.')
+    return cast('Triple', out)
+
+
+@dataclass(frozen=True, slots=True)
+class IRSparseSupport:
+    """Sparse relation/support attributes for an IR operation."""
+
+    kind: str
+    kernel_size: Triple | None = None
+    stride: Triple | None = None
+    padding: Triple | None = None
+    dilation: Triple | None = None
+    target: str | None = None
+    mode: str | None = None
+    join: str | None = None
+
+
+@dataclass(frozen=True, slots=True)
+class IRTensorSpec:
+    """Named graph input or output specification."""
+
+    name: str
+    type: IRValueType
+
+
+@dataclass(frozen=True, slots=True)
+class IRNode:
+    """One semantic operation in a lattice model graph."""
+
+    id: str
+    op: str
+    inputs: dict[str, IRInputRef]
+    outputs: dict[str, str]
+    parameters: dict[str, IRParameter] = field(default_factory=dict)
+    attributes: dict[str, Any] = field(default_factory=dict)
+    support: IRSparseSupport | None = None
+
+
+@dataclass(frozen=True, slots=True)
+class IRManifest:
+    """Validated sparse model manifest.
+
+    The manifest is the stable artifact contract shared by future training
+    producers and the MLX artifact loader. It records semantic sparse graph
+    nodes and names the tensor weights stored beside it.
+    """
+
+    schema_version: str
+    producer: dict[str, str]
+    runtime: dict[str, str]
+    inputs: tuple[IRTensorSpec, ...]
+    outputs: tuple[IRTensorSpec, ...]
+    nodes: tuple[IRNode, ...]
+    dtype_policy: DTypePolicy = 'preserve'
+    coordinate_order: tuple[str, ...] = _COORDINATE_ORDER
+    feature_layout: tuple[str, ...] = _FEATURE_LAYOUT
+    weight_layout: str = _WEIGHT_LAYOUT
+
+
+def load_manifest(path: str | Path) -> IRManifest:
+    """Load and validate a manifest JSON file."""
+
+    with Path(path).open('r', encoding='utf-8') as file:
+        raw = json.load(file)
+    return manifest_from_dict(raw)
+
+
+def manifest_from_dict(raw: Mapping[str, Any]) -> IRManifest:
+    """Build a validated :class:`IRManifest` from decoded JSON data."""
+
+    _require_mapping(raw, 'manifest')
+    schema_version = _require_str(raw, 'schema_version')
+    if schema_version != CURRENT_SCHEMA_VERSION:
+        raise ValueError(
+            f'unsupported lattice IR schema_version {schema_version!r}; '
+            f'expected {CURRENT_SCHEMA_VERSION!r}.'
+        )
+
+    coordinate_order = _str_tuple(
+        raw.get('coordinate_order', _COORDINATE_ORDER),
+        'coordinate_order',
+    )
+    if coordinate_order != _COORDINATE_ORDER:
+        raise ValueError(
+            "coordinate_order must be ['batch', 'x', 'y', 'z']."
+        )
+
+    feature_layout = _str_tuple(
+        raw.get('feature_layout', _FEATURE_LAYOUT),
+        'feature_layout',
+    )
+    if feature_layout != _FEATURE_LAYOUT:
+        raise ValueError("feature_layout must be ['row', 'channel'].")
+
+    weight_layout = _require_str(
+        raw, 'weight_layout', default=_WEIGHT_LAYOUT
+    )
+    if weight_layout != _WEIGHT_LAYOUT:
+        raise ValueError("weight_layout must be 'mlx-lattice'.")
+
+    dtype_policy = _dtype_policy(
+        _require_str(raw, 'dtype_policy', default='preserve')
+    )
+
+    manifest = IRManifest(
+        schema_version=schema_version,
+        producer=_str_map(raw.get('producer', {}), 'producer'),
+        runtime=_str_map(raw.get('runtime', {}), 'runtime'),
+        coordinate_order=coordinate_order,
+        feature_layout=feature_layout,
+        weight_layout=weight_layout,
+        dtype_policy=dtype_policy,
+        inputs=tuple(
+            _tensor_spec(item, f'inputs[{index}]')
+            for index, item in enumerate(_require_list(raw, 'inputs'))
+        ),
+        outputs=tuple(
+            _tensor_spec(item, f'outputs[{index}]')
+            for index, item in enumerate(_require_list(raw, 'outputs'))
+        ),
+        nodes=tuple(
+            _node(item, f'nodes[{index}]')
+            for index, item in enumerate(_require_list(raw, 'nodes'))
+        ),
+    )
+    _validate_graph(manifest)
+    return manifest
+
+
+def manifest_to_dict(manifest: IRManifest) -> dict[str, Any]:
+    """Convert a manifest object to JSON-serializable data."""
+
+    return {
+        'schema_version': manifest.schema_version,
+        'producer': dict(manifest.producer),
+        'runtime': dict(manifest.runtime),
+        'coordinate_order': list(manifest.coordinate_order),
+        'feature_layout': list(manifest.feature_layout),
+        'weight_layout': manifest.weight_layout,
+        'dtype_policy': manifest.dtype_policy,
+        'inputs': [
+            {'name': item.name, 'type': item.type}
+            for item in manifest.inputs
+        ],
+        'outputs': [
+            {'name': item.name, 'type': item.type}
+            for item in manifest.outputs
+        ],
+        'nodes': [_node_to_dict(node) for node in manifest.nodes],
+    }
+
+
+def is_ir_value_type(value: str) -> bool:
+    """Return whether ``value`` is a supported lattice IR value type."""
+
+    return value in _IR_VALUE_TYPES
+
+
+def ir_value_type(value: str) -> IRValueType:
+    """Validate and cast a string into :class:`IRValueType`."""
+
+    if not is_ir_value_type(value):
+        raise ValueError(f'unsupported IR value type: {value!r}.')
+    return cast('IRValueType', value)
+
+
+def _node_to_dict(node: IRNode) -> dict[str, Any]:
+    raw: dict[str, Any] = {
+        'id': node.id,
+        'op': node.op,
+        'inputs': {
+            key: value if isinstance(value, str) else list(value)
+            for key, value in node.inputs.items()
+        },
+        'outputs': dict(node.outputs),
+    }
+    if node.parameters:
+        raw['parameters'] = dict(node.parameters)
+    if node.attributes:
+        raw['attributes'] = dict(node.attributes)
+    if node.support is not None:
+        raw['support'] = _support_to_dict(node.support)
+    return raw
+
+
+def _support_to_dict(support: IRSparseSupport) -> dict[str, Any]:
+    raw: dict[str, Any] = {'kind': support.kind}
+    for name in ('kernel_size', 'stride', 'padding', 'dilation'):
+        value = getattr(support, name)
+        if value is not None:
+            raw[name] = list(value)
+    for name in ('target', 'mode', 'join'):
+        value = getattr(support, name)
+        if value is not None:
+            raw[name] = value
+    return raw
+
+
+def _support_attributes(support: IRSparseSupport) -> dict[str, Any]:
+    raw: dict[str, Any] = {}
+    for name in ('kernel_size', 'stride', 'padding', 'dilation'):
+        value = getattr(support, name)
+        if value is not None:
+            raw[name] = value
+    if support.target is not None:
+        raw['coordinates'] = support.target
+    if support.mode is not None:
+        raw['mode'] = support.mode
+    if support.join is not None:
+        raw['join'] = support.join
+    return raw
+
+
+def _node(raw: Any, path: str) -> IRNode:
+    data = _require_mapping(raw, path)
+    attributes = dict(
+        _require_mapping(
+            data.get('attributes', {}),
+            f'{path}.attributes',
+        )
+    )
+    support = (
+        None
+        if 'support' not in data
+        else _support(data['support'], f'{path}.support')
+    )
+    if support is not None:
+        attributes.update(_support_attributes(support))
+    return IRNode(
+        id=_require_str(data, 'id', path=path),
+        op=_require_str(data, 'op', path=path),
+        inputs=_input_map(data.get('inputs', {}), f'{path}.inputs'),
+        outputs=_str_map(data.get('outputs', {}), f'{path}.outputs'),
+        parameters=_str_map(
+            data.get('parameters', {}), f'{path}.parameters'
+        ),
+        attributes=attributes,
+        support=support,
+    )
+
+
+def _support(raw: Any, path: str) -> IRSparseSupport:
+    data = _require_mapping(raw, path)
+    return IRSparseSupport(
+        kind=_require_str(data, 'kind', path=path),
+        kernel_size=_optional_triple(data, 'kernel_size', path),
+        stride=_optional_triple(data, 'stride', path),
+        padding=_optional_triple(data, 'padding', path),
+        dilation=_optional_triple(data, 'dilation', path),
+        target=_optional_str(data, 'target', path),
+        mode=_optional_str(data, 'mode', path),
+        join=_optional_str(data, 'join', path),
+    )
+
+
+def _tensor_spec(raw: Any, path: str) -> IRTensorSpec:
+    data = _require_mapping(raw, path)
+    name = _require_str(data, 'name', path=path)
+    value_type = _require_str(data, 'type', path=path)
+    try:
+        typed = ir_value_type(value_type)
+    except ValueError as exc:
+        raise ValueError(
+            f'{path}.type is not a supported lattice IR value type.'
+        ) from exc
+    return IRTensorSpec(name, typed)
+
+
+def _validate_graph(manifest: IRManifest) -> None:
+    names: set[str] = set()
+    node_ids: set[str] = set()
+    for item in manifest.inputs:
+        _require_unique(names, item.name, 'graph value')
+    for node in manifest.nodes:
+        _require_unique(node_ids, node.id, 'node id')
+        if not node.outputs:
+            raise ValueError(f'node {node.id!r} must define outputs.')
+        for value_ref in node.inputs.values():
+            for value_name in _value_refs(value_ref):
+                if value_name not in names:
+                    raise ValueError(
+                        f'node {node.id!r} references unknown input '
+                        f'{value_name!r}.'
+                    )
+        for value_name in node.outputs.values():
+            _require_unique(names, value_name, 'graph value')
+    for item in manifest.outputs:
+        if item.name not in names:
+            raise ValueError(f'unknown manifest output {item.name!r}.')
+
+
+def _require_unique(seen: set[str], value: str, label: str) -> None:
+    if value in seen:
+        raise ValueError(f'duplicate {label}: {value!r}.')
+    seen.add(value)
+
+
+def _optional_triple(
+    data: Mapping[str, Any],
+    name: str,
+    path: str,
+) -> Triple | None:
+    if name not in data:
+        return None
+    return triple(data[name], name=f'{path}.{name}')
+
+
+def _dtype_policy(value: str) -> DTypePolicy:
+    if value not in ('preserve', 'fp32', 'fp16', 'fp16_inference'):
+        raise ValueError(
+            'dtype_policy must be preserve, fp32, fp16, or fp16_inference.'
+        )
+    return cast('DTypePolicy', value)
+
+
+def _str_tuple(raw: Any, path: str) -> tuple[str, ...]:
+    if not isinstance(raw, list | tuple):
+        raise ValueError(f'{path} must be a list of strings.')
+    values = tuple(raw)
+    if not all(isinstance(item, str) for item in values):
+        raise ValueError(f'{path} must be a list of strings.')
+    return cast(tuple[str, ...], values)
+
+
+def _str_map(raw: Any, path: str) -> dict[str, str]:
+    data = _require_mapping(raw, path)
+    out: dict[str, str] = {}
+    for key, value in data.items():
+        if not isinstance(key, str) or not isinstance(value, str):
+            raise ValueError(f'{path} must map strings to strings.')
+        out[key] = value
+    return out
+
+
+def _input_map(raw: Any, path: str) -> dict[str, IRInputRef]:
+    data = _require_mapping(raw, path)
+    out: dict[str, IRInputRef] = {}
+    for key, value in data.items():
+        if not isinstance(key, str):
+            raise ValueError(f'{path} must map strings to value refs.')
+        if isinstance(value, str):
+            out[key] = value
+            continue
+        if isinstance(value, list | tuple) and all(
+            isinstance(item, str) for item in value
+        ):
+            out[key] = tuple(value)
+            continue
+        raise ValueError(
+            f'{path}.{key} must be a string or list of strings.'
+        )
+    return out
+
+
+def _value_refs(value: IRInputRef) -> tuple[str, ...]:
+    return (value,) if isinstance(value, str) else value
+
+
+def _require_mapping(raw: Any, path: str) -> Mapping[str, Any]:
+    if not isinstance(raw, Mapping):
+        raise ValueError(f'{path} must be an object.')
+    return raw
+
+
+def _require_list(raw: Mapping[str, Any], name: str) -> Sequence[Any]:
+    value = raw.get(name)
+    if not isinstance(value, list):
+        raise ValueError(f'{name} must be a list.')
+    return value
+
+
+def _require_str(
+    raw: Mapping[str, Any],
+    name: str,
+    *,
+    path: str = '',
+    default: str | None = None,
+) -> str:
+    value = raw.get(name, default)
+    if not isinstance(value, str):
+        prefix = f'{path}.' if path else ''
+        raise ValueError(f'{prefix}{name} must be a string.')
+    return value
+
+
+def _optional_str(
+    raw: Mapping[str, Any],
+    name: str,
+    path: str,
+) -> str | None:
+    if name not in raw:
+        return None
+    value = raw[name]
+    if not isinstance(value, str):
+        raise ValueError(f'{path}.{name} must be a string.')
+    return value

lattice_contract-0.1.0/lattice_contract/ops.py

@@ -0,0 +1,210 @@
+from __future__ import annotations
+
+from collections.abc import Callable, Iterator, Mapping
+from dataclasses import dataclass, field
+from typing import Literal, TypeVar, cast
+
+from lattice_contract.manifest import IRInputRef, IRNode, IRValueType
+
+DeclarationT = TypeVar('DeclarationT', bound=Callable)
+FunctionT = TypeVar('FunctionT', bound=Callable)
+IRParameterKind = Literal[
+    'array',
+    'optional_array',
+    'quantized_weight',
+    'array_or_quantized_weight',
+]
+
+
+@dataclass(frozen=True, slots=True)
+class IROpSpec:
+    """Static semantic contract for one lattice IR operation."""
+
+    name: str
+    inputs: frozenset[str]
+    outputs: frozenset[str]
+    output_types: dict[str, IRValueType]
+    input_types: dict[str, IRValueType] = field(default_factory=dict)
+    value_attribute_types: dict[str, IRValueType] = field(
+        default_factory=dict
+    )
+    parameters: frozenset[str] = frozenset()
+    optional_parameters: frozenset[str] = frozenset()
+    attributes: frozenset[str] = frozenset()
+    value_attributes: frozenset[str] = frozenset()
+    requires_support: bool = False
+
+
+@dataclass(frozen=True, slots=True)
+class IROpArtifactHints:
+    """Artifact hints that cannot be inferred from annotations alone."""
+
+    parameters: Mapping[str, IRParameterKind] = field(default_factory=dict)
+    optional_parameters: Mapping[str, IRParameterKind] = field(
+        default_factory=dict
+    )
+    attributes: frozenset[str] = frozenset()
+    value_attributes: frozenset[str] = frozenset()
+
+
+_OP_SPECS: dict[str, IROpSpec] = {}
+_ARTIFACT_HINT_ATTR = '__mlx_lattice_op_artifact_hints__'
+
+
+def lattice_op_hints(
+    *,
+    parameters: Mapping[str, str] | None = None,
+    optional_parameters: Mapping[str, str] | None = None,
+    attributes: set[str] | None = None,
+    value_attributes: set[str] | None = None,
+) -> Callable[[FunctionT], FunctionT]:
+    """Attach artifact classification hints to a public op function.
+
+    Most operation bindings are inferred from annotations. Hints are reserved
+    for ambiguous tensor arguments, especially persisted weights and optional
+    graph-carried values whose annotations include more than one IR type.
+    """
+
+    hints = IROpArtifactHints(
+        parameters=_parameter_kinds(parameters),
+        optional_parameters=_parameter_kinds(optional_parameters),
+        attributes=frozenset(attributes or ()),
+        value_attributes=frozenset(value_attributes or ()),
+    )
+
+    def decorator(function: FunctionT) -> FunctionT:
+        setattr(function, _ARTIFACT_HINT_ATTR, hints)
+        return function
+
+    return decorator
+
+
+def op_artifact_hints(function: Callable) -> IROpArtifactHints:
+    """Return artifact hints attached by :func:`lattice_op_hints`."""
+
+    return cast(
+        IROpArtifactHints,
+        getattr(function, _ARTIFACT_HINT_ATTR, IROpArtifactHints()),
+    )
+
+
+def _parameter_kinds(
+    values: Mapping[str, str] | None,
+) -> dict[str, IRParameterKind]:
+    return {
+        name: _parameter_kind(value)
+        for name, value in dict(values or {}).items()
+    }
+
+
+def _parameter_kind(value: str) -> IRParameterKind:
+    if value not in (
+        'array',
+        'optional_array',
+        'quantized_weight',
+        'array_or_quantized_weight',
+    ):
+        raise ValueError(f'unsupported IR parameter kind: {value!r}.')
+    return cast('IRParameterKind', value)
+
+
+def ir_op_spec(
+    name: str,
+    *,
+    inputs: set[str],
+    outputs: set[str],
+    output_types: Mapping[str, IRValueType] | None = None,
+    input_types: Mapping[str, IRValueType] | None = None,
+    value_attribute_types: Mapping[str, IRValueType] | None = None,
+    parameters: set[str] | None = None,
+    optional_parameters: set[str] | None = None,
+    attributes: set[str] | None = None,
+    value_attributes: set[str] | None = None,
+    requires_support: bool = False,
+) -> Callable[[DeclarationT], DeclarationT]:
+    """Register the semantic contract for one IR operation.
+
+    The annotation keeps the stable IR operation set compact and discoverable
+    without maintaining a hand-written registry table. Artifact consumers
+    attach their implementations separately, so importing
+    :mod:`lattice_contract` exposes the complete semantic contract without
+    importing an MLX, Torch, or native graph executor.
+    """
+
+    spec = IROpSpec(
+        name=name,
+        inputs=frozenset(inputs),
+        outputs=frozenset(outputs),
+        output_types=dict(output_types or {}),
+        input_types=dict(input_types or {}),
+        value_attribute_types=dict(value_attribute_types or {}),
+        parameters=frozenset(parameters or ()),
+        optional_parameters=frozenset(optional_parameters or ()),
+        attributes=frozenset(attributes or ()),
+        value_attributes=frozenset(value_attributes or ()),
+        requires_support=requires_support,
+    )
+
+    def decorator(declaration: DeclarationT) -> DeclarationT:
+        if name in _OP_SPECS:
+            raise ValueError(
+                f'duplicate lattice IR op registration: {name}.'
+            )
+        _OP_SPECS[name] = spec
+        return declaration
+
+    return decorator
+
+
+def iter_op_specs() -> Iterator[IROpSpec]:
+    """Iterate registered IR operation specs."""
+
+    return iter(_OP_SPECS.values())
+
+
+def op_spec(name: str) -> IROpSpec:
+    """Return the registered spec for ``name`` or raise ``ValueError``."""
+
+    try:
+        return _OP_SPECS[name]
+    except KeyError:
+        pass
+    raise ValueError(f'unsupported lattice IR op: {name!r}.')
+
+
+def validate_node_against_spec(node: IRNode) -> None:
+    """Validate node ports/parameters against the registered op spec."""
+
+    spec = op_spec(node.op)
+    _require_keys(node.inputs, spec.inputs, f'{node.id}.inputs')
+    _require_keys(node.outputs, spec.outputs, f'{node.id}.outputs')
+    allowed = spec.parameters | spec.optional_parameters
+    missing = spec.parameters - set(node.parameters)
+    extra = set(node.parameters) - allowed
+    if missing:
+        raise ValueError(
+            f'{node.id}.parameters missing required keys: '
+            f'{sorted(missing)}.'
+        )
+    if extra:
+        raise ValueError(
+            f'{node.id}.parameters has unsupported keys: {sorted(extra)}.'
+        )
+    if spec.requires_support and node.support is None:
+        raise ValueError(f'{node.id} requires a support object.')
+
+
+def _require_keys(
+    values: Mapping[str, IRInputRef],
+    expected: frozenset[str],
+    path: str,
+) -> None:
+    actual = set(values)
+    missing = expected - actual
+    extra = actual - expected
+    if missing:
+        raise ValueError(
+            f'{path} missing required keys: {sorted(missing)}.'
+        )
+    if extra:
+        raise ValueError(f'{path} has unsupported keys: {sorted(extra)}.')

lattice_contract-0.1.0/pyproject.toml

@@ -0,0 +1,32 @@
+[project]
+name = "lattice-contract"
+version = "0.1.0"
+description = "Backend-neutral sparse lattice artifact contract"
+readme = "README.md"
+requires-python = ">=3.12"
+license = "MIT"
+authors = [{ name = "Z.Y. Lin", email = "me@iki.moe" }]
+keywords = ["sparse", "lattice", "ir", "model-artifact"]
+
+[build-system]
+requires = ["hatchling>=1.27"]
+build-backend = "hatchling.build"
+
+[dependency-groups]
+dev = ["pytest>=8", "ruff>=0.15.15", "ty>=0.0.42"]
+
+[tool.hatch.build.targets.wheel]
+packages = ["lattice_contract"]
+
+[tool.ruff]
+line-length = 76
+target-version = "py312"
+
+[tool.ruff.lint]
+select = ["E4", "E7", "E9", "F", "I", "B", "UP", "SIM", "RUF"]
+fixable = ["ALL"]
+
+[tool.ruff.format]
+quote-style = "single"
+indent-style = "space"
+line-ending = "lf"

lattice_contract-0.1.0/tests/test_contract_manifest.py

@@ -0,0 +1,77 @@
+from __future__ import annotations
+
+import pytest
+
+from lattice_contract import (
+    ir_value_type,
+    is_ir_value_type,
+    manifest_from_dict,
+    manifest_to_dict,
+)
+
+
+def _manifest() -> dict:
+    return {
+        'schema_version': '0.1',
+        'producer': {'name': 'test'},
+        'runtime': {'name': 'mlx-lattice', 'version': '>=0.2,<0.3'},
+        'coordinate_order': ['batch', 'x', 'y', 'z'],
+        'feature_layout': ['row', 'channel'],
+        'weight_layout': 'mlx-lattice',
+        'dtype_policy': 'preserve',
+        'inputs': [{'name': 'input', 'type': 'sparse_tensor'}],
+        'outputs': [{'name': 'output', 'type': 'sparse_tensor'}],
+        'nodes': [
+            {
+                'id': 'relu',
+                'op': 'feature.relu',
+                'inputs': {'input': 'input'},
+                'outputs': {'output': 'output'},
+            }
+        ],
+    }
+
+
+def test_manifest_roundtrip_preserves_semantic_contract() -> None:
+    manifest = manifest_from_dict(_manifest())
+
+    raw = manifest_to_dict(manifest)
+
+    assert raw['schema_version'] == '0.1'
+    assert raw['coordinate_order'] == ['batch', 'x', 'y', 'z']
+    assert raw['feature_layout'] == ['row', 'channel']
+    assert raw['weight_layout'] == 'mlx-lattice'
+    assert raw['nodes'][0]['op'] == 'feature.relu'
+
+
+def test_manifest_rejects_unknown_schema_version() -> None:
+    raw = _manifest()
+    raw['schema_version'] = '9.9'
+
+    with pytest.raises(ValueError, match='unsupported lattice IR'):
+        manifest_from_dict(raw)
+
+
+def test_ir_value_type_helper_validates_schema_values() -> None:
+    assert is_ir_value_type('sparse_tensor')
+    assert ir_value_type('dense_tensor') == 'dense_tensor'
+    assert not is_ir_value_type('not_a_value_type')
+
+    with pytest.raises(ValueError, match='unsupported IR value type'):
+        ir_value_type('not_a_value_type')
+
+
+def test_manifest_rejects_unknown_graph_input_reference() -> None:
+    raw = _manifest()
+    raw['nodes'][0]['inputs']['input'] = 'missing'
+
+    with pytest.raises(ValueError, match='unknown input'):
+        manifest_from_dict(raw)
+
+
+def test_manifest_rejects_duplicate_graph_values() -> None:
+    raw = _manifest()
+    raw['nodes'][0]['outputs']['output'] = 'input'
+
+    with pytest.raises(ValueError, match='duplicate graph value'):
+        manifest_from_dict(raw)

lattice_contract-0.1.0/tests/test_contract_surface.py

@@ -0,0 +1,17 @@
+from __future__ import annotations
+
+import importlib
+import sys
+
+
+def test_lattice_contract_imports_without_mlx_lattice_runtime() -> None:
+    sys.modules.pop('lattice_contract', None)
+    sys.modules.pop('lattice_contract.manifest', None)
+    sys.modules.pop('lattice_contract.ops', None)
+    sys.modules.pop('mlx_lattice', None)
+
+    contract = importlib.import_module('lattice_contract')
+
+    assert contract.__version__ == '0.1.0'
+    assert contract.CURRENT_SCHEMA_VERSION == '0.1'
+    assert 'mlx_lattice' not in sys.modules