Metadata-Version: 2.4
Name: leapp
Version: 0.5.1
Summary: LEAPP - Lightweight Export Annotations for Policy Pipelines: A Python package for tracing and exporting computational graphs from PyTorch code
Author-email: Frank Lai <frlai@nvidia.com>
Maintainer-email: Frank Lai <frlai@nvidia.com>
License-Expression: Apache-2.0
Project-URL: Homepage, https://github.com/nvidia-isaac/leapp
Project-URL: Repository, https://github.com/nvidia-isaac/leapp
Project-URL: Issues, https://github.com/nvidia-isaac/leapp/issues
Keywords: machine-learning,pytorch,robotics,policy-pipelines,graph-export
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Developers
Classifier: Operating System :: OS Independent
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.8
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.8
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: torch>=2.6.0
Requires-Dist: matplotlib>=3.5.0
Requires-Dist: networkx>=2.6
Requires-Dist: PyYAML>=6.0
Requires-Dist: onnx>=1.19.0
Requires-Dist: onnxruntime>=1.20.0
Requires-Dist: onnxscript>=0.1.0
Requires-Dist: safetensors>=0.5.0
Provides-Extra: dev
Requires-Dist: pytest; extra == "dev"
Requires-Dist: black; extra == "dev"
Requires-Dist: flake8; extra == "dev"
Requires-Dist: mypy; extra == "dev"
Dynamic: license-file

<!--
SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
SPDX-License-Identifier: Apache-2.0
-->

# LEAPP - Lightweight Export Annotations for Policy Pipelines

A Python package for tracing and exporting multi-step PyTorch computational graphs. Annotate your existing code with lightweight markers, and LEAPP captures the graph structure, exports each stage as an individual model, and generates a deployment-ready YAML specification.

## What is LEAPP?

**LEAPP** stands for **Lightweight Export Annotations for Policy Pipelines**. It is designed for pipelines that chain multiple PyTorch models or processing stages together — where you need to export the whole pipeline, not just a single model.

LEAPP works by tracing real execution of your code. It records which tensors flow between stages, exports each stage independently, and writes a YAML that describes how to wire them back together at inference time. Unlike other solutions, LEAPP is designed to be non-invasive and eliminates the need to have another export-specific implementation of your processing.

## Features

- **Export to Multiple Formats**: TorchScript and ONNX
- **Flexible Structuring**: Easily define complex node boundaries or multi-node graph structures
- **Lightweight**: Minimal insertions — no rewrites of existing model code required, annotations safely noop if not exporting
- **YAML Specification**: Complete metadata for deployment and downstream frameworks
- **BYOM (Bring Your Own Model)**: Integrate pre-compiled models into the graph without recompiling
- **Automatic Graph Visualization**: Generate a diagram of your pipeline

## Installation

pip:
```
pip install leapp
```

## Documentation

For detailed guides and API reference, see the `docs/` directory.

## Usage



LEAPP operates using traced datatypes. Traced datatypes extend the original datatype and silently add tracing capabilities to them. The result is the ability to define node boundaries explicitly, including across helper functions or when inputs come from multiple call sites.

```python
import torch
import leapp
import env # mock environment setup
from leapp import annotate

def get_obs1_from_env():
    obs1 = env.get_from_env('obs')
    obs1 = annotate.input_tensors("preprocess", {"obs1_name": obs1}) #inserted annotation
    # **some processing**
    return obs1

def get_obs2_from_env():
    obs2 = env.get_from_env('obs2')
    obs2 = annotate.input_tensors("preprocess", {"obs2_name": obs2}) #inserted annotation
    # **some processing**
    return obs2


leapp.start(name="my_pipeline")

obs1 = get_obs1_from_env()
obs2 = get_obs2_from_env()

# Simple feature construction from the two observation tensors.
obs1_centered = obs1 - obs1.mean()
obs2_scaled = obs2 / (obs2.abs().max() + 1e-6)
features = torch.cat([obs1_centered, obs2_scaled], dim=-1)

annotate.output_tensors("preprocess", {"processed_features": features}, export_with="jit")  #inserted annotation

features = annotate.input_tensors("predict", {"features": features}) #inserted annotation
output = torch.relu(featues @ torch.randn(16, 4))
annotate.output_tensors("predict", {"output": output}, export_with="onnx") #inserted annotation

leapp.stop()
leapp.compile_graph()
```

On top of simply tracing inputs and outputs, leapp provides various annotations to accommodate a variety of policies. Please refer to `docs/api.md` for more details. 

## Output Files

Running `compile_graph()` creates a directory named after your graph containing:

- **`{graph_name}.yaml`** — complete pipeline specification (node I/O, model paths, data flow)
- **`{graph_name}.png`** — visual diagram of the pipeline
- **Individual model files** — one exported model per node (`.pt`, `.onnx`, etc...)

## API Reference

```python
import leapp
from leapp import annotate
```

**Lifecycle**

- `leapp.start(name, save_path=".", verbose=False, dry_run=False, non_traced=None, max_cached_io=5, global_patching=True)`
- `leapp.stop()`
- `leapp.compile_graph(visualize=True, verbose=None, validate=True, dry_run=False, rtol=1e-3, atol=1e-5, strict=True)`

**Annotations**

- `annotate.input_tensors(node_name, tensors)` — begin a traced tensor node
- `annotate.output_tensors(node_name, tensors, export_with, ...)` — finalize a traced tensor node
- `annotate.method(export_with, node_name, ...)` — decorator, shorthand for annotating i/o with function structure
- `annotate.state_tensors(node_name, tensors)` — declare recurrent/state inputs
- `annotate.update_state(node_name, tensors)` — update recurrent state outputs
- `annotate.register_buffer(node_name, tensors)` — register traced persistent buffers for in-place writes
- `annotate.module(node_name, model, ...)` — trace an `nn.Module` with buffer tracking
- `annotate.mirror_leapp_tags(source, target)` — copy tracing tags between tensors

See `docs/api.md` for the full reference and `docs/` for advanced features.

## Requirements

- Python ≥ 3.8
- PyTorch ≥ 2.6.0
- PyYAML ≥ 6.0
- matplotlib ≥ 3.5.0
- networkx ≥ 2.6
- onnx ≥ 1.19.0
- onnxruntime ≥ 1.20.0
- onnxscript ≥ 0.1.0
- safetensors ≥ 0.5.0

## License

Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.

Licensed under the Apache License, Version 2.0.
