Skip to content

Loader

loader

State-dict loading primitives: SDOps + safetensors loader.

SafetensorsStateDictLoader

Load state-dicts from one or more safetensors shards.

Uses :func:mlx.core.load for the actual tensor loading (which memory-maps the file). The :class:SDOps chain runs over each raw key — keys it drops are silently skipped, keys it rewrites end up in the output dict under their new name, and the optional key/value transforms run after the rename step.

The loader is stateless; create one and reuse it.

load 

load(path: str | Path | list[str | Path], sd_ops: SDOps | None = None) -> StateDict

Load one or more safetensors shards into a :class:StateDict.

Parameters:

Name Type Description Default
path str | Path | list[str | Path]

Single path or list of paths to merge. Later shards with overlapping keys overwrite earlier ones (matches dict.update semantics).

 required
sd_ops SDOps | None

Optional rename/match/transform chain.

 None

Returns:

Type Description
StateDict

The merged :class:StateDict. size accumulates each
StateDict

tensor's nbytes; dtype collects every distinct
StateDict

dtype seen.

ContentMatching dataclass 

ContentMatching(prefix: str = '', suffix: str = '')

Prefix/suffix matcher gating which keys :class:SDOps lets through.

Defaults to prefix="" and suffix="" which matches everything. A key passes :meth:SDOps.apply_to_key only if it matches at least one :class:ContentMatching in the chain.

ContentReplacement dataclass 

ContentReplacement(content: str, replacement: str)

A substring replacement applied to a state-dict key.

KeyValueOperation

Bases: Protocol

Callable protocol for arbitrary key/value transforms.

KeyValueOperationResult

Bases: NamedTuple

Output of a key-value transform: a (new key, new value) pair.

Operations may return multiple results (e.g. splitting a fused QKV tensor produces three results).

SDKeyValueOperation dataclass 

SDKeyValueOperation(key_matcher: ContentMatching, kv_operation: KeyValueOperation)

A :class:KeyValueOperation gated by a :class:ContentMatching.

SDOps dataclass 

SDOps(name: str, mapping: tuple[ContentReplacement | ContentMatching | SDKeyValueOperation, ...] = (), allowed_keys: frozenset[str] | None = None)

Immutable chain of state-dict operations.

Every fluent method returns a new :class:SDOps; the original is never mutated. This makes chains safe to share across loaders and safe to compose into reusable presets.

Attributes:

Name Type Description
name str

Human-readable label for the chain (used for debugging and serialization).
mapping tuple[ContentReplacement | ContentMatching | SDKeyValueOperation, ...]

Ordered tuple of operations to apply.
allowed_keys frozenset[str] | None

If set, keys not in this frozenset (after all replacements have been applied) are dropped.

with_replacement 

with_replacement(content: str, replacement: str) -> SDOps

Append a substring replacement to the chain.

with_matching 

with_matching(prefix: str = '', suffix: str = '') -> SDOps

Append a prefix/suffix filter to the chain.

Defaults to prefix="" / suffix="", which matches every key. Use empty defaults when you want to allow every key through and rely on later replacements / allowed-keys to gate them.

with_additional_allowed_keys 

with_additional_allowed_keys(keys: frozenset[str]) -> SDOps

Restrict (or further restrict) the post-replacement key set.

If :attr:allowed_keys is already set, the two sets are union'd — additive, never subtractive.

with_kv_operation 

with_kv_operation(operation: KeyValueOperation, key_prefix: str = '', key_suffix: str = '') -> SDOps

Append a key/value transform gated by an optional matcher.

apply_to_key 

apply_to_key(key: str) -> str | None

Apply matching + replacements. Returns None if the key is dropped.

A key is dropped when (a) no :class:ContentMatching in the chain accepts it, or (b) the post-replacement key is not in :attr:allowed_keys (when set).

apply_to_key_value 

apply_to_key_value(key: str, value: array) -> list[KeyValueOperationResult]

Apply the first matching :class:SDKeyValueOperation, if any.

Returns a single [(key, value)] if no kv operation matches — i.e. the value is passed through untouched. Only the first matching operation is invoked; design the chain accordingly.

StateDict dataclass 

StateDict(sd: dict[str, array], size: int, dtype: set[Dtype])

A loaded model state-dict with memory + dtype accounting.

Attributes:

Name Type Description
sd dict[str, array]

Parameter name → array.
size int

Total memory footprint in bytes (sum of array.nbytes).
dtype set[Dtype]

Set of distinct dtypes present (useful to detect mixed- precision shards).

footprint 

footprint() -> int

Total memory footprint in bytes.

read_safetensors_metadata 

read_safetensors_metadata(path: str | Path) -> dict[str, str]

Read the __metadata__ block from a safetensors file.

Safetensors layout: 8 little-endian bytes hold the JSON header length, then the header (UTF-8 JSON), then the raw tensor bytes. The optional __metadata__ key in the header carries arbitrary string-keyed string-valued model metadata (HF convention).

Parameters:

Name Type Description Default
path str | Path

Path to a .safetensors file.

 required

Returns:

Type Description
dict[str, str]

The metadata dict, or an empty dict if the file has no
dict[str, str]

__metadata__ block.

Raises:

Type Description
OSError

If the file cannot be read.
ValueError

If the header is malformed.

safetensors_loader

Safetensors state-dict loader with optional :class:SDOps key remapping.

SafetensorsStateDictLoader

Load state-dicts from one or more safetensors shards.

Uses :func:mlx.core.load for the actual tensor loading (which memory-maps the file). The :class:SDOps chain runs over each raw key — keys it drops are silently skipped, keys it rewrites end up in the output dict under their new name, and the optional key/value transforms run after the rename step.

The loader is stateless; create one and reuse it.

load 
load(path: str | Path | list[str | Path], sd_ops: SDOps | None = None) -> StateDict

Load one or more safetensors shards into a :class:StateDict.

Parameters:

Name Type Description Default
path str | Path | list[str | Path]

Single path or list of paths to merge. Later shards with overlapping keys overwrite earlier ones (matches dict.update semantics).

 required
sd_ops SDOps | None

Optional rename/match/transform chain.

 None

Returns:

Type Description
StateDict

The merged :class:StateDict. size accumulates each
StateDict

tensor's nbytes; dtype collects every distinct
StateDict

dtype seen.

read_safetensors_metadata 

read_safetensors_metadata(path: str | Path) -> dict[str, str]

Read the __metadata__ block from a safetensors file.

Safetensors layout: 8 little-endian bytes hold the JSON header length, then the header (UTF-8 JSON), then the raw tensor bytes. The optional __metadata__ key in the header carries arbitrary string-keyed string-valued model metadata (HF convention).

Parameters:

Name Type Description Default
path str | Path

Path to a .safetensors file.

 required

Returns:

Type Description
dict[str, str]

The metadata dict, or an empty dict if the file has no
dict[str, str]

__metadata__ block.

Raises:

Type Description
OSError

If the file cannot be read.
ValueError

If the header is malformed.

sd_ops

State-dict key renaming and matching operations.

A common need when porting models is to rewrite upstream weight-key naming (diffusers net.0.proj, ComfyUI diffusion_model.*, Megatron linear_fc1) into the local MLX module's parameter tree. :class:SDOps chains immutable rename / match / transform operations that are applied in order to each key (and optionally to each key/value pair) at load time.

Three operation kinds compose:

  • :class:ContentReplacement — substring rename (str.replace).
  • :class:ContentMatching — prefix/suffix filter; keys that don't match at least one such filter are dropped.
  • :class:SDKeyValueOperation — full key/value transform, e.g. to split a fused QKV tensor into three keys.

Build the chain with the fluent with_* methods on :class:SDOps, then pass the resulting :class:SDOps to a loader's load method::

ops = (
    SDOps("comfy_to_mlx")
    .with_matching()                            # accept all keys
    .with_replacement("diffusion_model.", "")
    .with_replacement(".to_out.0.", ".to_out.")
)

state = SafetensorsStateDictLoader().load(path, sd_ops=ops)

ContentReplacement dataclass 

ContentReplacement(content: str, replacement: str)

A substring replacement applied to a state-dict key.

ContentMatching dataclass 

ContentMatching(prefix: str = '', suffix: str = '')

Prefix/suffix matcher gating which keys :class:SDOps lets through.

Defaults to prefix="" and suffix="" which matches everything. A key passes :meth:SDOps.apply_to_key only if it matches at least one :class:ContentMatching in the chain.

KeyValueOperationResult

Bases: NamedTuple

Output of a key-value transform: a (new key, new value) pair.

Operations may return multiple results (e.g. splitting a fused QKV tensor produces three results).

KeyValueOperation

Bases: Protocol

Callable protocol for arbitrary key/value transforms.

SDKeyValueOperation dataclass 

SDKeyValueOperation(key_matcher: ContentMatching, kv_operation: KeyValueOperation)

A :class:KeyValueOperation gated by a :class:ContentMatching.

SDOps dataclass 

SDOps(name: str, mapping: tuple[ContentReplacement | ContentMatching | SDKeyValueOperation, ...] = (), allowed_keys: frozenset[str] | None = None)

Immutable chain of state-dict operations.

Every fluent method returns a new :class:SDOps; the original is never mutated. This makes chains safe to share across loaders and safe to compose into reusable presets.

Attributes:

Name Type Description
name str

Human-readable label for the chain (used for debugging and serialization).
mapping tuple[ContentReplacement | ContentMatching | SDKeyValueOperation, ...]

Ordered tuple of operations to apply.
allowed_keys frozenset[str] | None

If set, keys not in this frozenset (after all replacements have been applied) are dropped.
with_replacement 
with_replacement(content: str, replacement: str) -> SDOps

Append a substring replacement to the chain.

with_matching 
with_matching(prefix: str = '', suffix: str = '') -> SDOps

Append a prefix/suffix filter to the chain.

Defaults to prefix="" / suffix="", which matches every key. Use empty defaults when you want to allow every key through and rely on later replacements / allowed-keys to gate them.

with_additional_allowed_keys 
with_additional_allowed_keys(keys: frozenset[str]) -> SDOps

Restrict (or further restrict) the post-replacement key set.

If :attr:allowed_keys is already set, the two sets are union'd — additive, never subtractive.

with_kv_operation 
with_kv_operation(operation: KeyValueOperation, key_prefix: str = '', key_suffix: str = '') -> SDOps

Append a key/value transform gated by an optional matcher.

apply_to_key 
apply_to_key(key: str) -> str | None

Apply matching + replacements. Returns None if the key is dropped.

A key is dropped when (a) no :class:ContentMatching in the chain accepts it, or (b) the post-replacement key is not in :attr:allowed_keys (when set).

apply_to_key_value 
apply_to_key_value(key: str, value: array) -> list[KeyValueOperationResult]

Apply the first matching :class:SDKeyValueOperation, if any.

Returns a single [(key, value)] if no kv operation matches — i.e. the value is passed through untouched. Only the first matching operation is invoked; design the chain accordingly.

state_dict

State-dict container types.

StateDict dataclass 

StateDict(sd: dict[str, array], size: int, dtype: set[Dtype])

A loaded model state-dict with memory + dtype accounting.

Attributes:

Name Type Description
sd dict[str, array]

Parameter name → array.
size int

Total memory footprint in bytes (sum of array.nbytes).
dtype set[Dtype]

Set of distinct dtypes present (useful to detect mixed- precision shards).
footprint 
footprint() -> int

Total memory footprint in bytes.