Metadata-Version: 2.4
Name: torch_layer_segregator
Version: 0.1.0
Summary: Tag PyTorch submodules and export them as Torch, TorchScript, or Core ML models.
Requires-Python: >=3.10
Description-Content-Type: text/markdown
Requires-Dist: coremltools>=9.0
Dynamic: description
Dynamic: description-content-type
Dynamic: requires-dist
Dynamic: requires-python
Dynamic: summary

# torch_layer_segregator

Tag PyTorch submodules and export them as standalone Torch, TorchScript, or Core ML models.

Useful for profiling individual layers, benchmarking conversions, and building per-layer deployment pipelines without manually tracing each submodule.

## Install

```bash
pip install torch_layer_segregator
```

**Requirements:** Python 3.10+, PyTorch 2.7+. Core ML export also requires `coremltools` 9.0+ on macOS.

## Quick start

```python
import torch.nn as nn
from torch_layer_segregator import TorchLayerSaver, make_saveable_layer

class MyModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.block = make_saveable_layer(nn.Conv2d(3, 16, 3, padding=1), key="conv_block")

    def forward(self, x):
        return self.block(x)

model = MyModel()
input_shape = (1, 3, 224, 224)

saver = TorchLayerSaver(model, "exports/my_model", input_shape, device="cpu")
saver.save_torch_model_layers()
saver.save_script_model_layers()
saver.save_coreml_model_layers()
```

## How it works

1. **Tag layers** — wrap submodules with `make_saveable_layer` to mark them for export.
2. **Infer shapes** — `TorchLayerSaver` runs one forward pass and records each tagged layer's input and output shapes.
3. **Export** — save tagged layers to disk in the format you need.

Tagged layers are written under subdirectories of your output path, each with a `layer_details.json` manifest.

| Method | Directory | Output |
|---|---|---|
| `save_torch_model_layers()` | `torch_layers/` | `.pt` (metadata + `state_dict`) |
| `save_script_model_layers()` | `script_layers/` | TorchScript `.ts` |
| `save_coreml_model_layers()` | `coreml_layers/` | Core ML `.mlpackage` |

If no `key` is passed to `make_saveable_layer`, the filename is derived from the module's dotted name (e.g. `encoder.block0` → `encoder_block0`).

Core ML conversion traces each layer and exports an ML Program with float16 precision for all compute units (iOS 26+ deployment target).

## API

**`make_saveable_layer(module, key=None, input_shape=None)`**  
Mark a `nn.Module` for export. Returns the same module for inline use in `__init__`.

**`TorchLayerSaver(model, out_dir, input_shape, device="cpu")`**  
Run shape inference and coordinate exports. `input_shape` is the full model input tuple (batch included). `device` is used for inference and scripted/Core ML export.

All `save_*` methods accept an optional `out_dir` to override the default subdirectory.
