Metadata-Version: 2.4
Name: mlnode
Version: 1.0.0
Summary: JSON-described DAG → executable PyTorch nn.Module, with static shape validation
Project-URL: Homepage, https://github.com/Nasser-Obeid/MLnode
Project-URL: Repository, https://github.com/Nasser-Obeid/MLnode
Project-URL: Issues, https://github.com/Nasser-Obeid/MLnode/issues
Project-URL: Changelog, https://github.com/Nasser-Obeid/MLnode/blob/main/CHANGELOG.md
Project-URL: Documentation, https://github.com/Nasser-Obeid/MLnode/tree/main/docs
Project-URL: Editor, https://github.com/Nasser-Obeid/MLnode-editor
Author-email: Nasser Obeid <nasserobeid8@gmail.com>
License-Expression: MIT
License-File: LICENSE
Keywords: computational-graph,dag,deep-learning,json,model-builder,neural-networks,pytorch,visual-programming
Classifier: Development Status :: 4 - Beta
Classifier: Intended Audience :: Developers
Classifier: Intended Audience :: Science/Research
Classifier: Operating System :: OS Independent
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Programming Language :: Python :: 3.13
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Classifier: Typing :: Typed
Requires-Python: >=3.11
Requires-Dist: networkx>=3.0
Requires-Dist: pydantic>=2.0
Requires-Dist: torch>=2.4
Provides-Extra: dev
Requires-Dist: mypy>=1.10; extra == 'dev'
Requires-Dist: pytest-cov>=5.0; extra == 'dev'
Requires-Dist: pytest>=8.0; extra == 'dev'
Requires-Dist: ruff>=0.4; extra == 'dev'
Requires-Dist: transformers>=4.30; extra == 'dev'
Provides-Extra: hf
Requires-Dist: transformers>=4.30; extra == 'hf'
Description-Content-Type: text/markdown

# MLnode

**JSON-described DAG → executable PyTorch `nn.Module`.**

MLnode lets you define deep learning architectures as JSON graphs. It parses, validates, and compiles them into standard PyTorch modules you can train and deploy. No `exec()`, no `eval()`, no magic. MLnode speeds up ML/DL development by letting you focus on the model architecture instead of the wiring code — and by catching shape errors statically, before a single tensor is allocated.

Graphs can be written by hand, generated programmatically, or drawn visually in [mlnode-editor](https://github.com/Nasser-Obeid/MLnode-editor), the companion web editor that exports the same `.mlnode.json` format.

## Installation

```bash
pip install mlnode            # core library
pip install "mlnode[hf]"      # + HuggingFace nodes (transformers)
```

Requires Python ≥ 3.11. PyTorch ≥ 2.4, NetworkX ≥ 3.0, and Pydantic ≥ 2.0 are installed automatically.

For development, clone [the repo](https://github.com/Nasser-Obeid/MLnode) and `pip install -e ".[dev]"` — or run `./setup.sh`, which checks the Python version, creates `.venv`, and installs everything.

## Documentation

| Doc | Contents |
|-----|----------|
| [QUICKSTART.md](https://github.com/Nasser-Obeid/MLnode/blob/main/QUICKSTART.md) | Zero to trained model in five minutes |
| [docs/getting-started.md](https://github.com/Nasser-Obeid/MLnode/blob/main/docs/getting-started.md) | Install, first model, validate/build/save, training |
| [docs/json-schema.md](https://github.com/Nasser-Obeid/MLnode/blob/main/docs/json-schema.md) | Complete `.mlnode.json` format reference |
| [docs/api.md](https://github.com/Nasser-Obeid/MLnode/blob/main/docs/api.md) | Python API reference for every public class |
| [docs/validation.md](https://github.com/Nasser-Obeid/MLnode/blob/main/docs/validation.md) | How static shape validation works |
| [docs/blocks.md](https://github.com/Nasser-Obeid/MLnode/blob/main/docs/blocks.md) | Reusable sub-graphs: inline blocks, `block_ref`, `BlockRegistry` |
| [docs/custom-layers.md](https://github.com/Nasser-Obeid/MLnode/blob/main/docs/custom-layers.md) | Plugging your own `nn.Module` classes into a graph |
| [docs/huggingface.md](https://github.com/Nasser-Obeid/MLnode/blob/main/docs/huggingface.md) | Pretrained and config-initialized HuggingFace nodes |
| [docs/MLNODE_SPEC_v1_0.md](https://github.com/Nasser-Obeid/MLnode/blob/main/docs/MLNODE_SPEC_v1_0.md) | The frozen v1.0 specification |

## Quick Start

```python
from mlnode import MLnodeModel
import torch

model = MLnodeModel.from_dict({
    "mlnode_schema_version": "1.0",
    "input_shape": ["B", 10],
    "output_node": "fc",
    "nodes": [
        {"id": "input", "type": "Input", "params": {}},
        {"id": "fc", "type": "Linear", "params": {"in_features": 10, "out_features": 5}},
    ],
    "edges": [{"source": "input", "target": "fc"}]
})

result = model.validate()
print(result.valid)        # True
print(result.node_shapes)  # {'input': ('B', 10), 'fc': ('B', 5)}

module = model.build()
x = torch.randn(4, 10)
out = module(x)            # shape: (4, 5)
```

## Architecture

Every model goes through three distinct phases:

```
JSON → Parser → Validator → Executor → nn.Module
```

**Parser** takes JSON, validates it against the Pydantic schema, builds a NetworkX directed graph, runs topological sort, and detects cycles. Block nodes are expanded here by flattening their internal sub-graphs with namespaced IDs.

**Validator** propagates tensor shapes through the graph statically, without allocating GPU memory. Each layer type has a registered shape transform rule. Errors are collected (not fail-fast) so you can fix multiple issues in one pass.

**Executor** constructs a live `nn.Module` with proper weight sharing, reduction ops, and multi-output ports. The resulting module works like any PyTorch module — you can train it, save state dicts, export to ONNX, etc.

## API Reference

### MLnodeModel

The top-level class tying everything together.

```python
class MLnodeModel:
    @classmethod
    def load(cls, path: str, block_registry=None) -> MLnodeModel
    @classmethod
    def from_dict(cls, data: dict, block_registry=None) -> MLnodeModel
    def validate(self) -> ValidationResult
    def build(self) -> nn.Module
    def save(self, path: str, *, expanded: bool = False) -> None
```

**`load(path)`** — Load from a `.mlnode.json` file.

**`from_dict(data)`** — Create from a Python dictionary.

**`validate()`** — Run static shape validation. Returns a `ValidationResult` with `valid: bool`, `errors: list[ShapeError]`, and `node_shapes: dict[str, Shape]` mapping each node ID to its output shape.

**`build()`** — Compile to an executable `nn.Module`. Calls `validate()` first; raises `ValidationError` if invalid.

**`save(path)`** — Write to `.mlnode.json`. By default preserves block structure. Pass `expanded=True` to save the flattened graph.

### BlockRegistry

Stores reusable block definitions that can be referenced by `block_ref`.

```python
from mlnode import BlockRegistry, BlockDefinition

registry = BlockRegistry()
registry.register("resnet_block", BlockDefinition.model_validate({
    "input_node": "in", "output_node": "out",
    "nodes": [...], "edges": [...]
}))
registry.register_from_file("transformer_block", "blocks/transformer.json")

model = MLnodeModel.load("model.mlnode.json", block_registry=registry)
```

### Error Hierarchy

```
MLnodeError
├── ValidationError
│   ├── CycleDetectedError
│   └── ShapeMismatchError
├── CustomLayerError
└── SchemaVersionError
```

All errors carry context: `node_id`, `expected_shape`, `actual_shape`, `source_node_id` where applicable.

## JSON Schema Reference

### Top-Level Structure

```json
{
  "mlnode_schema_version": "1.0",
  "metadata": {"name": "My Model", "description": "...", "author": "..."},
  "input_shape": ["B", 3, 224, 224],
  "output_node": "classifier",
  "nodes": [...],
  "edges": [...]
}
```

**`mlnode_schema_version`** — Must be `"1.0"`. Unsupported versions raise `SchemaVersionError`.

**`metadata`** — Optional. Name, description, author.

**`input_shape`** — Shape fed to each Input node. Dimensions can be symbolic strings (`"B"`, `"Seq"`, `"D"`) or concrete integers. Symbolic dims let the validator check architecture correctness without fixing batch size or sequence length.

**`output_node`** — ID of the node whose output becomes the model's return value. Must exist in `nodes`.

### Node Schema

```json
{
  "id": "conv1",
  "type": "Conv2d",
  "params": {"in_channels": 3, "out_channels": 64, "kernel_size": 3, "padding": 1},
  "shared_id": "proj",
  "outputs": ["part1", "part2"],
  "reduction": "concat",
  "reduction_dim": -1,
  "hf_model": "bert-base-uncased",
  "hf_config": "GPT2Config",
  "path": "user.layers.MyLayer",
  "block": { ... },
  "block_ref": "resnet_block"
}
```

Only `id`, `type`, and `params` are required. All other fields are optional.

**`shared_id`** — Nodes with the same `shared_id` share a single `nn.Module` instance (weight sharing). They must have identical `type` and `params`.

**`outputs`** — Named output ports for multi-output nodes like `Split`. Referenced by edges via `source_port`.

**`reduction`** / **`reduction_dim`** — Fallback reduction strategy for non-reduction-op nodes that receive multiple inputs. Only allowed on nodes whose type is NOT a reduction op (`Add`, `Concat`, etc.).

**`hf_model`** — Load a pretrained HuggingFace model via `AutoModel.from_pretrained()`. Top-level field, not inside `params`. Requires the `[hf]` extra: `pip install "mlnode[hf]"`.

**`hf_config`** — Initialize from a HuggingFace config class (random weights, no download). Top-level field, not inside `params`. Requires the `[hf]` extra.

**`path`** — Import path for custom layers (e.g. `"user.layers.MyLayer"`). Only valid when `type` is `"custom"`. The class must inherit `nn.Module`.

**`block`** — Inline sub-graph definition for `Block` type nodes. Contains `input_node`, `output_node`, `nodes`, `edges`.

**`block_ref`** — Name of a block in the `BlockRegistry` for `Block` type nodes.

### Edge Schema

```json
{
  "source": "conv1",
  "source_port": "part1",
  "target": "fc1",
  "target_port": null
}
```

**`source_port`** — Selects a named output port when the source has multiple outputs. `null` for default single output.

**`target_port`** — Selects a named input port (e.g. `"query"`, `"key"`, `"value"` for attention). `null` for default.

### Supported Node Types

| Category | Types |
|----------|-------|
| I/O | `Input` |
| Linear | `Linear`, `Bilinear`, `LazyLinear`, `Identity` |
| Convolution | `Conv1d/2d/3d`, `ConvTranspose1d/2d/3d`, `LazyConv1d/2d/3d`*, `LazyConvTranspose1d/2d/3d`*, `Unfold`*, `Fold`* |
| Pooling | `MaxPool1d/2d/3d`, `AvgPool1d/2d/3d`, `AdaptiveAvgPool1d/2d/3d`, `AdaptiveMaxPool1d/2d/3d`, `MaxUnpool1d/2d/3d`*, `FractionalMaxPool2d/3d`*, `LPPool1d/2d/3d`* |
| Activation | `ReLU`, `ReLU6`, `GELU`, `SiLU`, `SELU`, `CELU`, `ELU`, `LeakyReLU`, `PReLU`, `RReLU`, `Sigmoid`, `Hardsigmoid`, `LogSigmoid`, `Tanh`, `Tanhshrink`, `Hardtanh`, `Hardswish`, `Hardshrink`, `Softplus`, `Softshrink`, `Softsign`, `Mish`, `Threshold`, `GLU`*, `Softmax`, `Softmin`, `Softmax2d`, `LogSoftmax`, `AdaptiveLogSoftmaxWithLoss`* |
| Normalization | `BatchNorm1d/2d/3d`, `LazyBatchNorm1d/2d/3d`, `LayerNorm`, `GroupNorm`, `SyncBatchNorm`, `InstanceNorm1d/2d/3d`, `LazyInstanceNorm1d/2d/3d`, `LocalResponseNorm`, `RMSNorm` |
| Dropout | `Dropout`, `Dropout1d/2d/3d`, `AlphaDropout`, `FeatureAlphaDropout` |
| Sparse | `Embedding`, `EmbeddingBag` |
| Recurrent | `RNN`, `LSTM`, `GRU`, `RNNCell`, `GRUCell`, `LSTMCell`* |
| Attention | `MultiheadAttention`, `Transformer`, `TransformerEncoder`, `TransformerDecoder`, `TransformerEncoderLayer`, `TransformerDecoderLayer` |
| Vision Ops | `PixelShuffle`, `PixelUnshuffle` |
| Reduction Ops | `Add`, `Subtract`, `Multiply`, `Divide`, `Concat`, `Stack` |
| Tensor Ops | `Flatten`, `Reshape`, `Permute`, `Transpose`, `View`, `Unsqueeze`, `Squeeze`, `Split` |
| Composable | `Block` |
| Custom | `custom` |

\* Accepted by the schema and buildable, but the static shape rule is not implemented yet — graphs containing these types currently fail `validate()` (and therefore `build()`). See [docs/validation.md](https://github.com/Nasser-Obeid/MLnode/blob/main/docs/validation.md) for the full status table.

Reduction ops have implicit reduction semantics — they accept multiple inputs without needing a `reduction` field. Setting `reduction` on a reduction op is an error.

**Multi-input layers.** Some layers take several positional tensor arguments instead of being reduced. Label their incoming edges with `target_port`:

| Type | Ports (in order) |
|------|------------------|
| `Transformer` | `src`, `tgt` |
| `TransformerDecoder`, `TransformerDecoderLayer` | `tgt`, `memory` |
| `MultiheadAttention` | `query`, `key`, `value` — one input means self-attention (`q=k=v`); two means `query` + shared `key`/`value` |
| `Bilinear` | `input1`, `input2` |

Unlabeled edges are assigned to ports in declaration order. `TransformerEncoder`/`TransformerDecoder` are built from flat layer params plus `num_layers` (default 6) and `norm` (bool — appends a final `LayerNorm`).

All transformer/attention/RNN shape rules assume `batch_first: true` — set it in params (the editor does this by default).

### Symbolic Dimensions

The shape validator understands these symbolic dimension names:

| Symbol | Meaning | Typical use |
|--------|---------|-------------|
| `B` | Batch size | First dimension of any tensor |
| `Seq` | Sequence length | NLP/sequence models |
| `C` | Channels | CNN feature maps |
| `H`, `W` | Height, Width | Spatial dimensions |
| `D` | Generic dimension | Embedding size, hidden size |
| `*` | Unknown/inferred | When shape can't be determined |

Symbolic dims propagate through the graph. `Conv2d` with input `(B, 3, 224, 224)` outputs `(B, 64, 112, 112)` — the `B` passes through unchanged.

## Custom Layers

Custom layers are Python classes that inherit `nn.Module`. MLnode validates them by:

1. Importing the class via `importlib` (no `exec`/`eval`)
2. Verifying it inherits `nn.Module`
3. Instantiating with the provided `params`
4. Running a dummy forward pass with `torch.zeros(2, *shape)` to probe the output shape
5. Caching the result

```json
{
  "id": "proj",
  "type": "custom",
  "path": "my_project.layers.BottleneckProjection",
  "params": {"in_features": 512, "bottleneck": 64}
}
```

```python
# my_project/layers.py
class BottleneckProjection(nn.Module):
    def __init__(self, in_features: int, bottleneck: int):
        super().__init__()
        self.down = nn.Linear(in_features, bottleneck)
        self.up = nn.Linear(bottleneck, in_features)

    def forward(self, x):
        return self.up(F.relu(self.down(x)))
```

## Blocks

Blocks are reusable sub-graphs. They're expanded at parse time — the rest of the system doesn't know they exist.

```json
{
  "id": "res1",
  "type": "Block",
  "block": {
    "input_node": "block_in",
    "output_node": "block_relu",
    "nodes": [
      {"id": "block_in", "type": "Input", "params": {}},
      {"id": "block_conv", "type": "Conv2d", "params": {"in_channels": 64, "out_channels": 64, "kernel_size": 3, "padding": 1}},
      {"id": "block_bn", "type": "BatchNorm2d", "params": {"num_features": 64}},
      {"id": "block_skip", "type": "Add", "params": {}},
      {"id": "block_relu", "type": "ReLU", "params": {}}
    ],
    "edges": [
      {"source": "block_in", "target": "block_conv"},
      {"source": "block_conv", "target": "block_bn"},
      {"source": "block_in", "target": "block_skip"},
      {"source": "block_bn", "target": "block_skip"},
      {"source": "block_skip", "target": "block_relu"}
    ]
  }
}
```

After expansion, internal nodes get namespaced IDs: `res1/block_conv`, `res1/block_bn`, etc. The block's input node is removed and incoming edges are rewired to the first internal nodes.

Blocks can be chained (`res1 → res2`), nested (blocks containing blocks, up to depth 10), and referenced from a registry via `block_ref` instead of inline `block`.

## Examples

Runnable training notebooks live in `examples/notebooks/`:

- `mlnode_mnist_training.ipynb` — CNN classifier on MNIST (`mnist_cnn.mlnode.json`)
- `mlnode_transformer_training.ipynb` — TransformerEncoder classifier over MNIST rows-as-sequence (`mnist_transformer.mlnode.json`)
- `mlnode_gan_training.ipynb` — DCGAN-style GAN where generator and discriminator are both MLnode graphs (`gan_generator.mlnode.json`, `gan_discriminator.mlnode.json`)

All `.mlnode.json` files open directly in mlnode-editor (Import → auto-layout).

### ResNet Skip Connection

```json
{
  "mlnode_schema_version": "1.0",
  "metadata": {"name": "ResNet Block"},
  "input_shape": ["B", 64, 56, 56],
  "output_node": "relu",
  "nodes": [
    {"id": "input", "type": "Input", "params": {}},
    {"id": "conv1", "type": "Conv2d", "params": {"in_channels": 64, "out_channels": 64, "kernel_size": 3, "padding": 1}},
    {"id": "conv2", "type": "Conv2d", "params": {"in_channels": 64, "out_channels": 64, "kernel_size": 3, "padding": 1}},
    {"id": "skip", "type": "Add", "params": {}},
    {"id": "relu", "type": "ReLU", "params": {}}
  ],
  "edges": [
    {"source": "input", "target": "conv1"},
    {"source": "conv1", "target": "conv2"},
    {"source": "input", "target": "skip"},
    {"source": "conv2", "target": "skip"},
    {"source": "skip", "target": "relu"}
  ]
}
```

### Split and Concat

```json
{
  "nodes": [
    {"id": "input", "type": "Input", "params": {}},
    {"id": "split", "type": "Split", "params": {"split_size_or_sections": [32, 32], "dim": -1}, "outputs": ["left", "right"]},
    {"id": "fc_left", "type": "Linear", "params": {"in_features": 32, "out_features": 16}},
    {"id": "fc_right", "type": "Linear", "params": {"in_features": 32, "out_features": 16}},
    {"id": "cat", "type": "Concat", "params": {"dim": -1}},
    {"id": "out", "type": "Linear", "params": {"in_features": 32, "out_features": 10}}
  ],
  "edges": [
    {"source": "input", "target": "split"},
    {"source": "split", "source_port": "left", "target": "fc_left"},
    {"source": "split", "source_port": "right", "target": "fc_right"},
    {"source": "fc_left", "target": "cat"},
    {"source": "fc_right", "target": "cat"},
    {"source": "cat", "target": "out"}
  ]
}
```

### Multi-Input Graph

```python
model = MLnodeModel.from_dict({
    "mlnode_schema_version": "1.0",
    "input_shape": ["B", 32],
    "output_node": "add",
    "nodes": [
        {"id": "enc_input", "type": "Input", "params": {}},
        {"id": "dec_input", "type": "Input", "params": {}},
        {"id": "fc1", "type": "Linear", "params": {"in_features": 32, "out_features": 16}},
        {"id": "fc2", "type": "Linear", "params": {"in_features": 32, "out_features": 16}},
        {"id": "add", "type": "Add", "params": {}}
    ],
    "edges": [
        {"source": "enc_input", "target": "fc1"},
        {"source": "dec_input", "target": "fc2"},
        {"source": "fc1", "target": "add"},
        {"source": "fc2", "target": "add"}
    ]
})
module = model.build()
out = module(torch.randn(4, 32), torch.randn(4, 32))  # two inputs
```

Multiple Input nodes map to `forward(*args)` by declaration order in the `nodes` array.

## Running Tests

```bash
pytest tests/ -v              # full suite (128 tests)
pytest tests/ --cov=mlnode    # with coverage (~88%)
pytest tests/unit/            # unit tests only
pytest tests/integration/     # full pipeline tests
pytest -m "not slow"          # skip tests that download HuggingFace models
```

Markers: `slow` (downloads HuggingFace models), `gpu` (requires CUDA).

## Project Structure

```
src/mlnode/
├── __init__.py          # Public API exports
├── model.py             # MLnodeModel — top-level API
├── errors.py            # Error hierarchy
├── schema/
│   ├── models.py        # Pydantic models (GraphSchema, NodeSchema, EdgeSchema, BlockDefinition)
│   └── types.py         # Dim, Shape, ValidationResult, constants
├── parser/
│   └── dag.py           # DAGParser → ParsedGraph (topo sort, cycle detection)
├── validator/
│   └── shapes.py        # ShapeValidator with per-type shape transform rules
├── executor/
│   └── builder.py       # ModuleBuilder → MLnodeModule(nn.Module)
├── layers/
│   ├── registry.py      # Maps type strings to nn.Module constructors
│   └── custom.py        # Custom layer validation (spec §4, steps 1-5)
├── blocks/
│   ├── registry.py      # BlockRegistry for reusable block definitions
│   └── expansion.py     # Flattens Block nodes into parent graph
├── huggingface/
│   └── loader.py        # Pretrained and config-based HF model loading
└── utils/
    └── __init__.py      # resolve_multi_inputs() — port binding shared by
                         # the executor and the shape validator
```

## Design Principles

1. **One schema, one parser** — JSON is the single source of truth.
2. **DAG only** — cycles rejected at parse time.
3. **Explicit over implicit** — all tensor flows are edges.
4. **Fail fast** — validation before instantiation.
5. **Versioned contract** — every JSON carries a schema version.
6. **No runtime codegen** — no `exec()`, `eval()`, `compile()`.
7. **Parser → Validator → Executor** — distinct phases with clear boundaries.
