optimaxx 0.2.0

An intelligent AI model optimization framework for PyTorch.

Optimaxx

Optimaxx is an intelligent AI model optimization framework for PyTorch. It automatically analyzes hardware, profiles models, recommends optimizations, applies safe transformations, benchmarks performance, and generates rich reports.

Vision

Developers should be able to optimize any PyTorch model with a single line of code:

import optimax

report = optimax.analyze(model)
optimized_model = optimax.optimize(model, goal="latency", hardware="auto")

Features

• Hardware Detection — Automatically detect CPU, GPU, CUDA, ROCm, Apple Metal, TPU, RAM, VRAM, and environment details.
• Model Analysis — Deep inspection of model architecture, parameters, FLOPs, memory usage, and execution graph.
• Profiling — Measure forward/backward latency, throughput, peak memory, and percentile statistics.
• Recommendation Engine — Generate actionable optimization suggestions with expected speedup, memory savings, risk, and compatibility.
• Optimization Engine — Modular, independent optimization passes (torch.compile, mixed precision, FlashAttention, quantization, operator fusion, gradient checkpointing, ONNX export, TensorRT) with validation and rollback.
• Benchmarking — Compare original vs. optimized models with statistical rigor.
• Rich Reports — JSON, HTML, Markdown, console, and PDF reports with charts and summaries.
• CLI — Full command-line interface for analysis, optimization, benchmarking, and diagnostics.
• Configuration — Support for pyproject.toml, YAML, JSON, and environment variables.
• Dry Run Mode — Preview what optimizations would be applied without touching the model.
• Explainability — Every recommendation explains why it helps, why it is safe, and when it may not help.
• Graph Visualization — ASCII, HTML, Markdown, and JSON visualizations of the model execution graph.
• Lazy Imports — Heavy dependencies (torch, numpy, psutil) are imported only when needed for fast startup.

Quick Start

pip install optimaxx

import optimax

model = torch.nn.TransformerEncoderLayer(d_model=512, nhead=8)
report = optimax.analyze(model)
optimized = optimax.optimize(model, goal="latency", hardware="auto")

CLI Usage

# Analyze a model file
optimaxx analyze model.py

# Benchmark a model
optimaxx benchmark --model model.py --input-shape 1 3 224 224

# Optimize a model
optimaxx optimize --model model.py --goal latency --output optimized.pt

# Preview optimizations without applying them
optimaxx optimize --model model.py --goal latency --dry-run

# Generate a report
optimaxx report --input report.json --format html

# Run diagnostics
optimaxx doctor

Documentation

Full documentation is available at https://optimax-ai.github.io/optimaxx.

Installation

pip install optimaxx

For development:

pip install optimaxx[dev]

For optional backends (ONNX, TensorRT):

pip install optimaxx[all]

API Reference

Core Functions

optimax.analyze(model)

Analyze a PyTorch model and return a detailed profile including parameters, layers, FLOPs, and memory estimates.

import torch
import optimax

model = torch.nn.Linear(10, 5)
profile = optimax.analyze(model)
print(profile.total_parameters)
print(profile.model_size_mb)

optimax.optimize(model, goal="latency", hardware="auto", ...)

Optimize a PyTorch model with the intelligent pipeline. Supports goals latency, memory, throughput, balanced. Passes can be auto-selected or explicitly provided.

import torch
import optimax

model = torch.nn.TransformerEncoderLayer(d_model=512, nhead=8)
input_fn = lambda: torch.randn(1, 10, 512)
result = optimax.optimize(
    model,
    goal="latency",
    input_fn=input_fn,
    return_full_result=True,
)
print(result.speedup)
print(result.applied_passes)

optimax.benchmark(model, input_fn, device, ...)

Benchmark a model and return latency, throughput, and memory statistics.

import torch
import optimax

model = torch.nn.Linear(10, 5)
input_fn = lambda: torch.randn(2, 10)
result = optimax.benchmark(model, input_fn, device="cpu")
print(result.profile.forward_latency.mean_ms)
print(result.profile.throughput)

optimax.detect_hardware()

Detect the current hardware environment (CPU, GPU, memory, software versions).

import optimax

hw = optimax.detect_hardware()
print(hw.cpu.brand)
print(hw.has_cuda)
print(hw.memory.total_mb)

optimax.generate_recommendations(model, profile, hardware)

Generate a list of optimization recommendations with explainability fields.

import torch
import optimax

model = torch.nn.Linear(10, 5)
profile = optimax.analyze(model)
hw = optimax.detect_hardware()
recs = optimax.generate_recommendations(model, profile, hw)
for rec in recs:
    print(rec.name, rec.expected_speedup)
    print(rec.why_it_helps)
    print(rec.why_it_is_safe)
    print(rec.why_it_may_not_help)

optimax.generate_report(...)

Create a comprehensive report from analysis, recommendations, benchmarks, and optimization results.

import optimax

report = optimax.generate_report(
    hardware=optimax.detect_hardware(),
    profile=optimax.analyze(model),
    recommendations=recs,
)
print(report.to_json())
print(report.to_html())
print(report.to_markdown())
print(report.to_console())

optimax.score_optimization(profile, hardware, benchmark)

Calculate a comprehensive optimization score (0-100) for a model.

import optimax

score = optimax.score_optimization(profile, hardware, benchmark)
print(score.overall)
print(score.suggestions)

optimax.detect_bottlenecks(profile, hardware)

Detect performance bottlenecks and root causes.

import optimax

bottlenecks = optimax.detect_bottlenecks(profile, hardware)
for b in bottlenecks:
    print(b.bottleneck_type, b.severity, b.recommendation)

optimax.detect_future_features(model, profile, hardware)

Detect future-looking feature compatibility and explain availability.

import optimax

features = optimax.detect_future_features(model, profile, hardware)
for f in features:
    print(f.name, f.available, f.description)

Configuration

optimax.load_config(path=None, env_prefix="OPTIMAX_")

Load configuration from files and environment variables.

import optimax

cfg = optimax.load_config()
print(cfg.goal)
print(cfg.benchmark_runs)

Plugins

optimax.register_pass(name, pass_class)

Register a custom optimization pass.

import optimax
from optimax.optimizer.base import OptimizationPass, PassResult

class MyPass(OptimizationPass):
    name = "my_pass"
    def apply(self, model, **kwargs):
        return PassResult(success=True, pass_name=self.name, message="Applied")

optimax.register_pass("my_pass", MyPass)

optimax.list_passes()

List all registered optimization passes.

import optimax

print(optimax.list_passes())

CLI Reference

optimax analyze <model_path>

Analyze a PyTorch model and print a report.

Option	Description
--output, -o	Output report file path
--format, -f	Output format: console, json, html, markdown
--verbose, -v	Verbose output
--debug	Debug mode

optimax benchmark <model_path>

Benchmark a PyTorch model.

Option	Default	Description
--input-shape, -s	1 3 224 224	Input shape as space-separated integers
--device, -d	auto	Device: cpu, cuda, auto
--runs, -r	10	Number of benchmark runs
--warmup, -w	3	Number of warmup runs
--backward, -b	False	Include backward pass
--verbose, -v	False	Verbose output

optimax optimize <model_path>

Optimize a PyTorch model and save the result.

Option	Default	Description
--goal, -g	latency	Optimization goal: latency, memory, throughput, balanced
--output, -o	required	Output optimized model path
--pass, -p	auto	Optimization passes to apply (repeatable)
--input-shape, -s	1 3 224 224	Input shape for validation
--device, -d	auto	Device for validation
--verbose, -v	False	Verbose output
--return-result	False	Return full OptimizationResult as JSON
--dry-run	False	Preview what would be applied without applying

optimax report <input_path>

Render a report from a JSON file.

Option	Default	Description
--format, -f	console	Output format: console, json, html, markdown
--output, -o	None	Output file path

optimax advisor <model_path>

Analyze a model and generate optimization recommendations.

Option	Default	Description
--format, -f	console	Output format: console, json, markdown
--verbose, -v	False	Verbose output

optimax auto-benchmark <model_path>

Automatically benchmark multiple configurations and generate a ranking table.

Option	Default	Description
--input-shape, -s	1 3 224 224	Input shape
--device, -d	auto	Device
--runs, -r	10	Number of runs
--warmup, -w	3	Number of warmup runs
--verbose, -v	False	Verbose output

optimax score <model_path>

Calculate a comprehensive optimization score (0-100) for a model.

Option	Default	Description
--input-shape, -s	1 3 224 224	Input shape
--device, -d	auto	Device
--verbose, -v	False	Verbose output

optimax bottleneck <model_path>

Detect performance bottlenecks and root causes.

Option	Default	Description
--input-shape, -s	1 3 224 224	Input shape
--device, -d	auto	Device
--verbose, -v	False	Verbose output

optimax future-features <model_path>

Detect future-looking feature compatibility and explain availability.

Option	Default	Description
--verbose, -v	False	Verbose output

optimax profile <profile_name> <model_path>

Apply a built-in optimization profile to a model.

Option	Default	Description
--output, -o	required	Output optimized model path
--input-shape, -s	1 3 224 224	Input shape
--device, -d	auto	Device
--verbose, -v	False	Verbose output

Built-in profiles: training, inference, edge_device, gpu_server, cpu_only, llm, vision, diffusion.

optimax doctor

Run diagnostics and print environment info.

Configuration Guide

Optimaxx supports configuration via pyproject.toml, YAML, JSON, and environment variables. Resolution order (later overrides earlier):

1. Defaults
2. pyproject.toml in current directory
3. optimax.yaml or optimax.json in current directory
4. Explicit file path
5. Environment variables

pyproject.toml

[tool.optimax]
goal = "latency"
benchmark_runs = 20
warmup_runs = 5
verbose = true
precision = "fp16"

YAML (optimax.yaml)

goal: memory
benchmark_runs: 15
warmup_runs: 3
verbose: false
report_formats:
  - console
  - json

JSON (optimax.json)

{
  "goal": "throughput",
  "benchmark_runs": 25,
  "warmup_runs": 5,
  "precision": "bf16"
}

Environment Variables

All environment variables use the OPTIMAX_ prefix:

export OPTIMAX_GOAL=latency
export OPTIMAX_BENCHMARK_RUNS=20
export OPTIMAX_WARMUP_RUNS=5
export OPTIMAX_VERBOSE=true
export OPTIMAX_PRECISION=fp16

Valid boolean values: true, 1, yes, false, 0, no.

Configuration Options

Option	Type	Default	Description
goal	string	latency	Optimization goal
hardware	string	auto	Hardware target
backend	string	None	Backend preference
precision	string	fp32	Default precision
passes	list	[]	Explicit pass names
verbose	bool	false	Verbose logging
debug	bool	false	Debug mode
report_formats	list	["console"]	Output formats
benchmark_runs	int	10	Number of benchmark runs
warmup_runs	int	3	Number of warmup runs
validation_tolerance	dict	{"rtol": 1e-3, "atol": 1e-4}	Validation tolerance
custom	dict	{}	Arbitrary custom values

Plugin Development Guide

Creating a Custom Optimization Pass

All optimization passes must inherit from OptimizationPass and implement apply(), check_compatibility(), and validate().

from optimax.optimizer.base import OptimizationPass, PassResult
from optimax.hardware import HardwareProfile
import torch

class MyCustomPass(OptimizationPass):
    """A custom optimization pass."""

    name = "my_custom_pass"
    description = "Applies a custom transformation to the model."

    def apply(self, model: torch.nn.Module, **kwargs) -> PassResult:
        """Apply the optimization pass.

        Returns:
            PassResult with success=True and optionally a modified model.
        """
        try:
            # Apply your transformation
            # model = transform(model)
            return PassResult(
                success=True,
                pass_name=self.name,
                message="Custom pass applied successfully.",
                model=model,
            )
        except Exception as e:
            return PassResult(
                success=False,
                pass_name=self.name,
                message=f"Custom pass failed: {e}",
            )

    def check_compatibility(self, model: torch.nn.Module, hardware: HardwareProfile) -> bool:
        """Check if the pass is compatible with the model and hardware."""
        return True

    def validate(self, model: torch.nn.Module, example_input: torch.Tensor) -> bool:
        """Validate that the optimized model produces correct outputs."""
        try:
            with torch.no_grad():
                output = model(example_input)
            return output is not None
        except Exception:
            return False

    def rollback(self, model: torch.nn.Module) -> torch.nn.Module:
        """Rollback the optimization if validation fails."""
        return model

Registering a Pass

import optimax
from optimax.optimizer.base import OptimizationPass, PassResult

class MyCustomPass(OptimizationPass):
    name = "my_custom_pass"
    def apply(self, model, **kwargs):
        return PassResult(success=True, pass_name=self.name, message="ok", model=model)
    def check_compatibility(self, model, hardware):
        return True
    def validate(self, model, example_input):
        return True
    def rollback(self, model):
        return model

optimax.register_pass("my_custom_pass", MyCustomPass)

# Now use it
result = optimax.optimize(model, passes=["my_custom_pass"])

PassResult Fields

Field	Type	Description
success	bool	Whether the pass succeeded
pass_name	str	Name of the pass
message	str	Human-readable message
model	nn.Module	Optional modified model
metrics	dict	Optional performance metrics

Safety Features Guide

Dry Run Mode

Preview what optimizations would be applied without modifying the model:

import optimax

result = optimax.optimize(model, input_fn=input_fn, dry_run=True, return_full_result=True)
print(result.messages)

CLI:

optimaxx optimize model.pt --goal latency --dry-run

Automatic Rollback

The intelligent pipeline automatically rolls back any pass that fails validation or does not improve the target metric. Each pass is benchmarked before and after; if the improvement is insufficient, the model is restored to its pre-pass state.

Numerical Validation

Every applied pass is validated against the original model using torch.allclose() with configurable tolerance (rtol and atol). If outputs deviate beyond tolerance, the pass is rolled back automatically.

result = optimax.optimize(
    model,
    input_fn=input_fn,
    tolerance={"rtol": 1e-3, "atol": 1e-4},
)

Deep Copy Protection

The pipeline never mutates the user's original model. All work is done on a deep copy, ensuring the original model remains safe regardless of pass failures.

original = model
optimized = optimax.optimize(model, input_fn=input_fn)
assert original is not optimized  # Original is untouched

Validation Toggle

Validation can be disabled for faster experimentation (not recommended for production):

result = optimax.optimize(model, input_fn=input_fn, validate=False)

Benchmark-Based Decision Making

The pipeline benchmarks the model before and after each pass. It only keeps passes that measurably improve the chosen goal (latency, memory, throughput, or balanced). This prevents theoretical optimizations that hurt real-world performance.

Contributing

We welcome contributions! Please read our Contributing Guide and Code of Conduct before submitting issues or pull requests.

License

Optimaxx is licensed under the Apache License 2.0.

Security

Please see our Security Policy for reporting vulnerabilities.