Write focused model architecture in .aur. Ship readable PyTorch.
Aurane is a compiler-oriented machine learning DSL with parsing, semantic checks, type/shape inference, IR lowering, optimization, profiling, visualization, and a polished CLI workflow.
Aurane keeps model authoring compact without hiding the generated code. It is designed for developers who want high-signal model definitions, fast feedback from analysis tools, and clean Python output that still feels familiar to PyTorch users.
flowchart LR
A[".aur source"] --> B["Parser"]
B --> C["Semantic + Type Checks"]
C --> D["Optimizer"]
D --> E["IR Lowering"]
E --> F["PyTorch Generator"]
F --> G["Readable Python"]
use torch
experiment SimpleExample:
seed = 42
device = "cpu"
model TinyNet:
input_shape = (3, 32, 32)
def forward(x):
x -> conv2d(16, kernel=3).relu
-> maxpool(2)
-> flatten()
-> dense(64).relu
-> dense(10)
Compile it:
aurane compile examples/simple.aur tiny_net.py --validate --formatGenerated shape:
class TinyNet(nn.Module):
def __init__(self):
super().__init__()
self.conv2d1 = nn.Conv2d(3, 16, 3, stride=1, padding=0)
self.dense1 = nn.Linear(3600, 64)
self.dense2 = nn.Linear(64, 10)
def forward(self, x):
x = F.relu(self.conv2d1(x))
x = F.max_pool2d(x, 2)
x = torch.flatten(x, 1)
x = F.relu(self.dense1(x))
x = self.dense2(x)
return xgit clone https://github.com/desenyon/aurane.git
cd aurane
# Compiler and CLI
pip install -e .
# Full ML extras
pip install -e ".[all]"
# Developer tools
pip install -e ".[dev]"Requirements:
- Python 3.10+
- Rich for the CLI experience
- PyTorch 2.0+ for running generated ML programs
| Command | Purpose | Example |
|---|---|---|
compile |
Generate Python from .aur |
aurane compile model.aur model.py --validate --format |
check |
Run semantic and type/shape checks | aurane check model.aur --semantic --types --json |
inspect |
Browse AST and model stats | aurane inspect model.aur --verbose --stats --export ast.json |
visualize |
Render architecture graphs | aurane visualize model.aur --format mermaid |
profile |
Estimate params, FLOPs, memory | aurane profile model.aur --detailed |
ir |
Dump lowered intermediate representation | aurane ir model.aur --format json |
watch |
Recompile on change | aurane watch model.aur model.py --analyze |
lint |
Find and fix simple source issues | aurane lint model.aur --auto-fix |
format |
Normalize Aurane source style | aurane format examples/ --check |
- Graph-aware parser: supports sequential chains and explicit graph-style forward definitions without confusing kwargs for assignments.
- Backend registry: code generation now routes through a backend layer instead of hard-coding one path.
- IR tooling: lower models into a structured intermediate representation with
aurane ir. - First-class checks: semantic analysis and type/shape validation are exposed through
aurane check. - Safer CLI exits:
python -m aurane.clinow preserves command return codes. - Stable cache keys: cached compilation output includes backend/options/schema information to avoid stale generated code.
- Better visualization: Mermaid and DOT output preserve labels, shapes, and parameter counts.
- Release-grade examples: the bundled CNN, ResNet, Transformer, GAN, and simple examples compile, check, and syntax-validate.
| Area | Supported |
|---|---|
| Models | model, input_shape, sequential forward chains, graph forward blocks |
| Layers | conv2d, dense, linear, flatten, maxpool, avgpool, dropout, batch_norm, batchnorm, reshape, embedding, multihead_attention, layer_norm, positional_encoding |
| Activations | relu, gelu, sigmoid, tanh, softmax, leaky_relu, residual |
| Analysis | semantic issues, type/shape inference, parameter counts, FLOPs estimates |
| Output | idiomatic PyTorch modules and training scaffolds |
The examples/ directory includes:
simple.aur- compact CNN startermnist.aur- MNIST training pipelineresnet.aur- deeper convolutional classifiertransformer.aur- language-model style architecturegan.aur- generator/discriminator pair
Run the full example smoke path:
for file in examples/*.aur; do
aurane check "$file" --semantic --types
aurane compile "$file" "/tmp/$(basename "$file" .aur).py" --quiet
doneaurane/
parser.py # .aur source to AST
semantic_analyzer.py # DSL-level diagnostics
type_checker.py # tensor shape/type checks
optimizer.py # AST optimization passes
ir.py # intermediate representation
backends/ # backend registry and torch backend
codegen_torch.py # PyTorch source generation
profiler.py # params/FLOPs/memory estimates
visualizer.py # rich, Mermaid, and DOT architecture output
cli/ # command-line interface
uv run --extra dev black --check aurane tests
uv run --extra dev mypy aurane
uv run --extra dev pytest -q
uv buildRelease smoke:
for file in examples/*.aur; do
out="/tmp/aurane-$(basename "$file" .aur).py"
uv run aurane compile "$file" "$out" --quiet
uv run aurane check "$file" --semantic --types --json >/tmp/aurane-check.json
doneAurane is released under the MIT License.