mlxsummary

Model inspection and summary tools for MLX neural networks on Apple Silicon.

Inspired from and similar to torchsummary for PyTorch, but designed specifically for MLX's module system.

Features

• 📊 Multiple output formats: Table, Tree, JSON, Markdown, Minimal
• 🔍 Detailed inspection: Layer paths, parameter counts, shapes
• 🎯 Filtering: Find layers by type, name pattern, or parameter count
• 📈 Statistics: Aggregate stats by layer type
• 🖥️ CLI support: Use from command line or Python
• 🧊 Freeze-aware: Track trainable vs frozen parameters

Installation

pip install mlxsummary

Or install from source:

git clone https://github.com/dhruvshr/mlxsummary.git
cd mlxsummary
pip install -e .

Quick Start

import mlx.nn as nn
from mlxsummary import summary

# Create a model
model = nn.Sequential(
    nn.Linear(784, 256),
    nn.ReLU(),
    nn.Dropout(0.1),
    nn.Linear(256, 128),
    nn.ReLU(),
    nn.Linear(128, 10),
)

# Print summary
summary(model)

Output:

=============================================================================================================================
 Model Summary: Sequential
=============================================================================================================================
Layer                                         Type                         Params Details           
-----------------------------------------------------------------------------------------------------------------------------
layers.0                                      Linear                      200,960 (784 → 256)       
layers.1                                      ReLU                              0                   
layers.2                                      Dropout                           0                   
layers.3                                      Linear                       32,896 (256 → 128)       
layers.4                                      ReLU                              0                   
layers.5                                      Linear                        1,290 (128 → 10)        
-----------------------------------------------------------------------------------------------------------------------------
Total Parameters:                                                         235,146
=============================================================================================================================

Output Formats

Table (default)

summary(model, format="table")

Tree

summary(model, format="tree")

📦 Sequential (235,146 params)
│   ├── 0: Linear (784 → 256) [200,960]
│   ├── 1: ReLU [0]
│   ├── 2: Dropout [0]
│   ├── 3: Linear (256 → 128) [32,896]
│   ├── 4: ReLU [0]
│   ├── 5: Linear (128 → 10) [1,290]

JSON

data = summary(model, format="json", print_output=False)

Markdown

summary(model, format="markdown")

Minimal (one-line)

summary(model, format="minimal")
# Output: Sequential: 235,146 params (6 layers)

Programmatic API

Inspector

For detailed programmatic access:

from mlxsummary import inspect

inspector = inspect(model)

# Get all layers
layers = inspector.get_layers()
for layer in layers:
    print(f"{layer.path}: {layer.total_params:,} params")

# Get statistics
stats = inspector.get_stats()
print(f"Total: {stats.total_params:,}")
print(f"Layer types: {stats.layer_type_counts}")

# Find specific layers
linear_layers = inspector.find_layers(layer_type=nn.Linear)
attention = inspector.find_layers(name_pattern="attention")
large_layers = inspector.find_layers(min_params=10000)

Convenience Functions

from mlxsummary import count_params, get_layers, get_stats, to_dict

# Count parameters
total = count_params(model)
trainable = count_params(model, trainable_only=True)

# Get layers
layers = get_layers(model)
linear_layers = get_layers(model, layer_type=nn.Linear)

# Get stats
stats = get_stats(model)

# Export to dict
data = to_dict(model)

Options

summary(
    model,
    format="table",           # Output format
    show_shapes=True,         # Show layer dimensions
    show_trainable=True,      # Show trainable params column
    show_frozen=False,        # Show frozen params column
    max_depth=None,           # Limit layer depth
    max_rows=None,            # Limit output rows
    include_zero_param=True,  # Include zero-param layers
    width=100,                # Output width
    print_output=True,        # Print vs return only
)

Command Line

# Summarize a model from a file
mlxsummary model.py

# Different formats
mlxsummary model.py --format tree
mlxsummary model.py --format json -o model.json
mlxsummary model.py --format markdown

# Options
mlxsummary model.py --max-depth 2
mlxsummary model.py --hide-zero
mlxsummary model.py --no-shapes

# Demo mode
mlxsummary --demo

Your model file should define a model variable or get_model() function:

# model.py
import mlx.nn as nn

model = nn.Sequential(
    nn.Linear(784, 256),
    nn.ReLU(),
    nn.Linear(256, 10),
)

API Reference

Classes

Class	Description
MLXInspector	Main inspector for detailed model analysis
LayerInfo	Information about a single layer
ModelStats	Aggregate statistics about a model
FormatterOptions	Options for output formatting
OutputFormat	Enum of available output formats

Functions

Function	Description
summary(model, ...)	Generate and print a model summary
inspect(model)	Create an inspector instance
count_params(model)	Count model parameters
get_layers(model)	Get list of layer information
get_stats(model)	Get aggregate statistics
to_dict(model)	Export model info to dictionary
tree(model)	Shortcut for tree format
table(model)	Shortcut for table format

Requirements

• macOS with Apple Silicon (M1/M2/M3/M4)
• Python 3.9+
• MLX 0.1.0+

License

MIT License - see LICENSE file for details.

Contributing

Contributions are welcome! Please open an issue or pull request.

See Also

• MLX Documentation
• MLX GitHub
• torchsummary - Similar tool for PyTorch