mlx-arsenal
Low-level operations and reusable building blocks missing from MLX core — the toolbox you want when porting PyTorch models to Apple Silicon.
Tip: if you use Claude Code for MLX ports, the
mlx-portingskill teaches Claude to reach formlx-arsenalsubmodules (diffusion,spatial,attention,norm,encoding,moe,tiling, etc.) before hand-rolling ops.
Install
pip install mlx-arsenal
Or directly from source:
pip install git+https://github.com/dgrauet/mlx-arsenal.git
Modules
| Module | Components | Replaces (PyTorch) |
|---|---|---|
mlx_arsenal.spatial |
interpolate_nearest, interpolate_3d, avg_pool1d, replicate_pad, upsample_nearest/bilinear, pixel_shuffle/unshuffle, patchify/unpatchify, PatchEmbed2d/3d |
F.interpolate, F.avg_pool1d, F.pad(mode="replicate"), F.pixel_shuffle |
mlx_arsenal.layout |
to_channels_last/first, channels_last ctx manager, convert_conv_weights, load_safetensors |
NCHW ↔ NHWC conversion, weight transposition |
mlx_arsenal.conv |
weight_norm, WeightNorm |
nn.utils.weight_norm |
mlx_arsenal.attention |
causal_mask, sliding_window_mask, spatial_only_mask, temporal_only_mask, sliding_tile_block_mask, sliding_tile_centered_mask, radial_box_mask, radial_gaussian_mask, frame_stride_diagonal_mask, vertical_stripe_mask, classify_heads_from_qk, classify_heads_from_probs, classify, Kind, block_contiguous_permutation, invert_permutation, block_causal_mask, centroid_compensated_attention, probe_residual_correction, select_probe_rows, tile_labels, antidiagonal_block_scores, top_p_block_mask, block_self_similarity, select_tiling, minmax_block_scores, top_k_block_mask, radius_bounded_block_scores, relative_block_mask |
Attention mask creation (LLM, block-diffusion LLM, sparse video DiT), head-pattern profiler, SVG2 block-contiguous token permutation, sparse-block compensation (SVG-EAR centroids, SparsePR probe repair), dynamic block masks (XAttention, SPADE adaptive tiling, RBS radius-bounded scores) |
mlx_arsenal.norm |
PixelNorm, ScaleNorm |
Custom normalization layers |
mlx_arsenal.encoding |
FourierEmbedder |
Sinusoidal positional encoding |
mlx_arsenal.diffusion |
get_timestep_embedding, TimestepEmbedding, get_sampling_sigmas, dynamic_shift_schedule, FlowMatchEulerDiscreteScheduler, DDIMScheduler, euler_step, CurvatureAdaptiveStepper, classifier_free_guidance, posterior_mean_capped_guidance, TeaCacheController, sea_filter, PerLayerAttentionCache, PerHeadAttentionCache, splice_heads, cfg_head_similarity, cfg_skip_mask, CFGSimilarityProfiler, CFGSkipController, WindowResidualController, VerifiedFeatureCache, geometric_threshold, token_stats, threshold_transfer, topk_transfer, factor_transfer, entropy_bound_transfer, transfer_schedule, block_ranges, pooled_qk, qk_drift, HeadMaskCache |
Flow-matching + DDIM diffusion primitives, CAT-Flow adaptive step sizes, PMC-CFG capped guidance, TeaCache + SeaCache spectral distance, AST attention caches, ASC cond/uncond skip, WA-RS residual sharing, masked diffusion LLM (dLLM) block-decoding commit rules, head-wise sparse-mask reuse (HEART) |
mlx_arsenal.moe |
MoEGate, MoELayer |
Top-k mixture-of-experts dispatch |
mlx_arsenal.rasterize |
rasterize_triangles, interpolate |
Tile-binned triangle rasterization with exact fixed-point coverage (Metal kernels) |
mlx_arsenal.tiling |
tiled_process, temporal_slice_process |
Memory-efficient large tensor processing |
mlx_arsenal.streaming |
BlockStreamer, BlockLoraSource, LoraFuser |
Low-RAM transformer block streaming from mmap'd safetensors |
mlx_arsenal.modulation |
AdaLNModulation, ScaleShiftTable, modulate, gated_residual |
DiT AdaLN modulation primitives (1 / 2 / 6 / 9-param variants) |
mlx_arsenal.ffn |
FeedForward, GatedFFN, GeGLU, SwiGLU |
Transformer FFN / MLP blocks (vanilla + gated variants) |
mlx_arsenal.loader |
SDOps, SafetensorsStateDictLoader, StateDict, read_safetensors_metadata |
State-dict key remapping chain + safetensors loader |
mlx_arsenal.rope |
rope_frequencies_1d, rope_frequencies_nd, apply_rotary_emb, rotate_half, meshgrid_nd |
Rotary Position Embeddings (N-D, interleaved + half-rotated variants) |
Quick start
from mlx_arsenal.spatial import interpolate_nearest, avg_pool1d, replicate_pad
from mlx_arsenal.layout import to_channels_last, convert_conv_weights
from mlx_arsenal.attention import causal_mask
# Resize a video tensor (B, D, H, W, C)
x_resized = interpolate_nearest(x, size=(8, 32, 32))
# Temporal pooling
pooled = avg_pool1d(temporal_features, kernel_size=2)
# Pad with edge replication (like F.pad mode="replicate")
padded = replicate_pad(x, [(0,0), (2,0), (1,1), (1,1), (0,0)])
# Convert PyTorch conv weights to MLX channels-last layout
mlx_weights = convert_conv_weights(pytorch_weights)
# Causal attention mask for autoregressive decoding
mask = causal_mask(seq_len=128, offset=kv_cache_len)
Block streaming (low-RAM transformers)
Run a 20+ GB transformer on a Mac without holding every block resident at once: keep one shared block module, and rebind its weights from memory-mapped safetensors before each block's forward.
from mlx_arsenal.streaming import BlockStreamer
# Build the model with ONE block in transformer_blocks (not num_layers).
model = build_my_transformer(num_layers=1)
load_non_block_weights(model, weights_path)
streamer = BlockStreamer(
weights_path,
block_prefix="transformer.transformer_blocks.",
)
assert streamer.block_count == num_layers # discovered from safetensors
shared_block = model.transformer_blocks[0]
prev_idx = None
for i in range(streamer.block_count):
streamer.bind(shared_block, idx=i, evict_previous=prev_idx)
x = shared_block(x, ...) # use the rebound block
prev_idx = i
For LoRA: pass a lora_fuser callable to BlockStreamer and one or
more BlockLoraSource instances to bind(..., lora_sources=...).
Quantization-aware fusion strategies stay in the caller — arsenal
only handles the discovery + indexing.
Requirements
- Python >= 3.10
- MLX >= 0.32.1
- Apple Silicon Mac
Development
pip install -e ".[dev]"
pytest tests/
# Optional: install the pre-commit hook so ruff runs on every `git commit`.
pip install pre-commit
pre-commit install
License
Apache 2.0
1 maintainer
Release files for mlx-arsenal 0.16.0
For a detailed explanation of source distributions (sdists) and built distributions (wheels), please see the package formats documentation.
Source distribution (sdist)
| File | Size | Uploaded | |
|---|---|---|---|
| mlx_arsenal-0.16.0.tar.gz | 188.3 kB | Details |
Built distribution (wheel)
| File | Interpreter | ABI | Platform | Reset |
|---|---|---|---|---|
| mlx_arsenal-0.16.0-py3-none-any.whl | Python 3 | none | any | Details |
Total release size: 317.0 kB
Release files / mlx_arsenal-0.16.0.tar.gz
| Download URL | mlx_arsenal-0.16.0.tar.gz |
|---|---|
| Size | 188.3 kB |
| Tags | Source |
|
SHA-256 checksum How to use checksums |
cd26b8c1f8467ed7f232c6544bd6fb0c275d162f6e7f40161a43bc41e9c9b3c2
|
|
BLAKE2b-256 checksum How to use checksums |
9a8bc53f0be009d881daed940f5664d8562942a8c0b01ce4edce788755ffcb7c
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
Yes |
| Uploaded via |
twine/7.0.0 CPython/3.13.14
|
Provenance
Provenance describes where a file came from. On PyPI, provenance is shared via attestations, which provide a verifiable record of the build or publishing details. View details, limitations and caveats.
PyPI Publish Attestation
PyPI verified that this artifact, at this checksum, originated from the publisher listed below.
Signed by GitHub Actions, verified by PyPI on Sep 29, 2026.
Transparency logRelease files / mlx_arsenal-0.16.0-py3-none-any.whl
| Download URL | mlx_arsenal-0.16.0-py3-none-any.whl |
|---|---|
| Size | 128.7 kB |
| Tags | Python 3 |
|
SHA-256 checksum How to use checksums |
a09e11dc796a084d938af659f8955bf1a0f5776bd5bd4afb1968504922929fa9
|
|
BLAKE2b-256 checksum How to use checksums |
e8078b89c6b762c082109f1c3647f3bb978e663d49278ddea06c55d0d15fa4a2
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
Yes |
| Uploaded via |
twine/7.0.0 CPython/3.13.14
|
Provenance
Provenance describes where a file came from. On PyPI, provenance is shared via attestations, which provide a verifiable record of the build or publishing details. View details, limitations and caveats.
PyPI Publish Attestation
PyPI verified that this artifact, at this checksum, originated from the publisher listed below.
Signed by GitHub Actions, verified by PyPI on Sep 29, 2026.
Transparency log
