Metadata-Version: 2.4
Name: spartan-torch
Version: 0.0.1
Summary: Custom layers and building blocks for PyTorch — assemble and fine-tune models fast
Keywords: pytorch,deep-learning,transformer,attention,vision-transformer,fine-tuning
Author: mievst
License-Expression: MIT
License-File: LICENSE
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Developers
Classifier: Intended Audience :: Science/Research
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Programming Language :: Python :: 3.13
Classifier: Programming Language :: Python :: 3.14
Classifier: Programming Language :: Python :: 3.15
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Requires-Dist: torch>=2.13.0
Requires-Dist: torchvision>=0.28.0
Requires-Dist: albumentationsx>=2.3.7 ; extra == 'dev'
Requires-Dist: ipywidgets>=8.1.8 ; extra == 'dev'
Requires-Dist: jupyterlab>=4.6.2 ; extra == 'dev'
Requires-Dist: matplotlib>=3.11.1 ; extra == 'dev'
Requires-Dist: pytest>=8.3.0 ; extra == 'dev'
Requires-Dist: tqdm>=4.70.0 ; extra == 'dev'
Requires-Dist: datasets>=3.0 ; extra == 'experiments'
Requires-Dist: ipywidgets>=8 ; extra == 'experiments'
Requires-Dist: lightning>=2.5 ; extra == 'experiments'
Requires-Dist: litlogger>=2026.6.25 ; extra == 'experiments'
Requires-Dist: mlflow>=2.20 ; extra == 'experiments'
Requires-Dist: timm>=1.0 ; extra == 'experiments'
Requires-Dist: tokenizers>=0.20 ; extra == 'experiments'
Requires-Dist: transformers[torch]>=4.50 ; extra == 'experiments'
Requires-Dist: trl>=1.0 ; extra == 'experiments'
Requires-Python: >=3.11
Project-URL: Homepage, https://github.com/mievst/spartan-torch
Project-URL: Repository, https://github.com/mievst/spartan-torch
Project-URL: Issues, https://github.com/mievst/spartan-torch/issues
Project-URL: Changelog, https://github.com/mievst/spartan-torch/blob/main/CHANGELOG.md
Provides-Extra: dev
Provides-Extra: experiments
Description-Content-Type: text/markdown

# spartan-torch

[![CI](https://github.com/mievst/spartan-torch/actions/workflows/ci.yml/badge.svg)](https://github.com/mievst/spartan-torch/actions/workflows/ci.yml)
[![PyPI](https://img.shields.io/pypi/v/spartan-torch.svg)](https://pypi.org/project/spartan-torch/)
[![Python](https://img.shields.io/pypi/pyversions/spartan-torch.svg)](https://pypi.org/project/spartan-torch/)

Custom layers and building blocks for PyTorch — assemble and fine-tune models fast.

## Quickstart

```python
import torch
from spartan_torch import TransformerBlock

block = TransformerBlock(
    in_size=128, head_size=32, num_heads=4,
    out_size=128, ff_hidden_size=512,
)
x = torch.randn(2, 16, 128)
y, _ = block(x)  # y: (2, 16, 128), cache kept for KV decoding
```

## What's inside

Drop-in primitives missing from `torch.nn`, one brick at a time:

- **`spartan_torch.transformers`** — attentions (MHA/GQA, linear, Linformer, Performer, Reformer), SSM mixers (Mamba / Mamba-2 / Mamba-3), pre/post/hybrid-norm blocks, norms (`RMSNorm`, …), positionals (RoPE, ALiBi, learned/sinusoidal), chunked and reversible variants;
- **`spartan_torch.cnn`** — residual, bottleneck, depthwise-separable and inverted-residual blocks;
- **`spartan_torch.vit`** — patch embeddings, class tokens, position embeddings, MAE decoder head;
- **`spartan_torch.masking` / `spartan_torch.selfsup`** — patch masking, DINO heads/loss, momentum encoder;
- **`spartan_torch.utils`** — small training helpers (e.g. warmup scheduler).

Two things this library deliberately does **not** do:

- **No full models.** Stock architectures (ResNet, ViT, LLaMA, …) come from `torchvision` / `timm` / 🤗 `transformers` with pretrained weights. `spartan_torch.compat` ships key-mappings that load those weights straight into our primitives.
- **No heavy training frameworks.** The core depends only on `torch` / `torchvision` — no Lightning, no HF `Trainer`. Experiment harnesses live in `experiments/` as optional extras.

## Reproducibility

- [`RESULTS.md`](https://github.com/mievst/spartan-torch/blob/main/RESULTS.md) — weight parity of our primitives against timm/torchvision references, plus a latency/memory sweep of sub-quadratic attentions with a Pareto plot. Tables are generated by scripts (`scripts/verify_pretrained.py`, `scripts/bench_attention.py`), never edited by hand.

## Installation

Requires Python **3.11+** (and a CUDA-capable Torch — see below).

```bash
pip install spartan-torch
```

For contributors, editable install with experiment and dev extras:

```bash
pip install -e ".[experiments,dev]"
```

> **CUDA wheels (uv):** in `pyproject.toml` (`[tool.uv.sources]` / `[[tool.uv.index]]`) torch is pinned to the `cu132` index (`download.pytorch.org/whl/cu132`). `pip` from PyPI installs the stock torch build; for a specific CUDA release install with an explicit `--index-url https://download.pytorch.org/whl/<cuda-version>`.

## Development

Dependencies and tests via [just](https://github.com/casey/just) + uv:

```bash
just sync   # install/sync the environment (uv sync --extra dev --extra experiments)
just test   # run the test suite (venv python -m pytest)
just parity # reference weight-parity tests (timm/torchvision; pretrained pulls weights over the network)
just verify # refresh the parity table in RESULTS.md
just bench  # GPU attention sweep + Pareto plot (requires CUDA)
just clean  # delete .venv
```

> Don't run `just sync` / `just clean` while Jupyter kernels are alive — they lock `.pyd`/`.so` files and uv can't finish the sync.
>
> Development happens in a devcontainer: `.devcontainer/` (CUDA image, GPU passthrough, host-mounted HuggingFace/W&B caches).

## License

MIT — see [LICENSE](https://github.com/mievst/spartan-torch/blob/main/LICENSE).
