optimaxx 0.1.0

An intelligent AI model optimization framework for PyTorch.

Optimax

Optimax 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.

Quick Start

pip install optimax

import torch
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
optimax analyze model.py

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

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

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

# Run diagnostics
optimax doctor

Documentation

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

Installation

pip install optimax

For development:

pip install optimax[dev]

For optional backends (ONNX, TensorRT):

pip install optimax[all]

Contributing

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

License

Optimax is licensed under the Apache License 2.0.

Security

Please see our Security Policy for reporting vulnerabilities.