Metadata-Version: 2.4
Name: optimaxx
Version: 0.2.0
Summary: An intelligent AI model optimization framework for PyTorch.
Project-URL: Homepage, https://github.com/ivanho-git/optimaxx
Project-URL: Repository, https://github.com/ivanho-git/optimaxx
Project-URL: Issues, https://github.com/ivanho-git/optimaxx/issues
Project-URL: Changelog, https://github.com/ivanho-git/optimaxx/blob/main/CHANGELOG.md
Project-URL: Documentation, https://ivanho-git.github.io/optimaxx
Author-email: Ibhan Mukherjee <ibhaanm29@email.com>
License: Apache-2.0
License-File: LICENSE
Keywords: compiler,deep-learning,inference,machine-learning,optimization,performance,profiler,pytorch,quantization
Classifier: Development Status :: 4 - Beta
Classifier: Intended Audience :: Developers
Classifier: Intended Audience :: Science/Research
Classifier: License :: OSI Approved :: Apache Software License
Classifier: Operating System :: OS Independent
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.9
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Classifier: Topic :: Software Development :: Libraries :: Python Modules
Requires-Python: >=3.9
Requires-Dist: jinja2>=3.1.0
Requires-Dist: numpy>=1.24.0
Requires-Dist: packaging>=23.0
Requires-Dist: psutil>=5.9.0
Requires-Dist: pydantic>=2.0.0
Requires-Dist: pyyaml>=6.0
Requires-Dist: rich>=13.0.0
Requires-Dist: toml>=0.10.2
Requires-Dist: torch>=2.0.0
Requires-Dist: typer>=0.9.0
Provides-Extra: all
Requires-Dist: onnx>=1.14.0; extra == 'all'
Requires-Dist: onnxruntime>=1.15.0; extra == 'all'
Requires-Dist: tensorrt>=8.6.0; (platform_machine == 'x86_64') and extra == 'all'
Provides-Extra: dev
Requires-Dist: black>=23.7.0; extra == 'dev'
Requires-Dist: mypy>=1.5.0; extra == 'dev'
Requires-Dist: pre-commit>=3.3.0; extra == 'dev'
Requires-Dist: pytest-cov>=4.1.0; extra == 'dev'
Requires-Dist: pytest-xdist>=3.3.0; extra == 'dev'
Requires-Dist: pytest>=7.4.0; extra == 'dev'
Requires-Dist: ruff>=0.1.0; extra == 'dev'
Provides-Extra: docs
Requires-Dist: mkdocs-material>=9.1.0; extra == 'docs'
Requires-Dist: mkdocs>=1.5.0; extra == 'docs'
Requires-Dist: mkdocstrings[python]>=0.22.0; extra == 'docs'
Description-Content-Type: text/markdown

# Optimaxx

[![CI](https://github.com/optimax-ai/optimaxx/actions/workflows/ci.yml/badge.svg)](https://github.com/optimax-ai/optimaxx/actions)
[![PyPI](https://img.shields.io/pypi/v/optimaxx)](https://pypi.org/project/optimaxx/)
[![Python](https://img.shields.io/pypi/pyversions/optimaxx)](https://pypi.org/project/optimaxx/)
[![License](https://img.shields.io/pypi/l/optimaxx)](https://github.com/optimax-ai/optimaxx/blob/main/LICENSE)
[![Coverage](https://img.shields.io/codecov/c/github/optimax-ai/optimaxx)](https://codecov.io/gh/optimax-ai/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:

```python
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

```bash
pip install optimaxx
```

```python
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

```bash
# 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](https://optimax-ai.github.io/optimaxx).

## Installation

```bash
pip install optimaxx
```

For development:

```bash
pip install optimaxx[dev]
```

For optional backends (ONNX, TensorRT):

```bash
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.

```python
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.

```python
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.

```python
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).

```python
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.

```python
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.

```python
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.

```python
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.

```python
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.

```python
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.

```python
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.

```python
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.

```python
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

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

### YAML (`optimax.yaml`)

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

### JSON (`optimax.json`)

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

### Environment Variables

All environment variables use the `OPTIMAX_` prefix:

```bash
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()`.

```python
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

```python
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:

```python
import optimax

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

CLI:
```bash
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.

```python
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.

```python
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):

```python
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](CONTRIBUTING.md) and [Code of Conduct](CODE_OF_CONDUCT.md) before submitting issues or pull requests.

## License

Optimaxx is licensed under the [Apache License 2.0](LICENSE).

## Security

Please see our [Security Policy](SECURITY.md) for reporting vulnerabilities.
