lmxlab 0.2.0__tar.gz

lmxlab-0.2.0/.github/dependabot.yml

@@ -0,0 +1,11 @@
+version: 2
+updates:
+  - package-ecosystem: "pip"
+    directory: "/"
+    schedule:
+      interval: "weekly"
+
+  - package-ecosystem: "github-actions"
+    directory: "/"
+    schedule:
+      interval: "weekly"

lmxlab-0.2.0/.github/pull_request_template.md

@@ -0,0 +1,13 @@
+## Summary
+
+<!-- 1-3 sentences describing what this PR does and why -->
+
+## Changes
+
+-
+
+## Checklist
+
+- [ ] Tests added/updated
+- [ ] Lint passes (`uv run ruff check src/ tests/ recipes/`)
+- [ ] Format passes (`uv run ruff format --check src/ tests/ recipes/`)

lmxlab-0.2.0/.github/workflows/ci.yml

@@ -0,0 +1,42 @@
+name: CI
+
+on:
+  push:
+    branches: [main]
+  pull_request:
+    branches: [main]
+
+jobs:
+  lint:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v6
+      - uses: astral-sh/setup-uv@v7
+        with:
+          enable-cache: true
+          cache-dependency-glob: "uv.lock"
+      - run: uv sync --locked --extra dev
+      - run: uv run ruff check src/ tests/ recipes/
+      - run: uv run ruff format --check src/ tests/ recipes/
+
+  docs:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v6
+      - uses: astral-sh/setup-uv@v7
+        with:
+          enable-cache: true
+          cache-dependency-glob: "uv.lock"
+      - run: uv sync --locked --extra dev
+      - run: uv run mkdocs build --strict
+
+  test:
+    runs-on: macos-14
+    steps:
+      - uses: actions/checkout@v6
+      - uses: astral-sh/setup-uv@v7
+        with:
+          enable-cache: true
+          cache-dependency-glob: "uv.lock"
+      - run: uv sync --locked --extra dev
+      - run: uv run pytest tests/ -v -m "not slow"

lmxlab-0.2.0/.github/workflows/dependabot-auto-merge.yml

@@ -0,0 +1,27 @@
+name: Dependabot Auto-Merge
+
+on: pull_request
+
+permissions:
+  contents: write
+  pull-requests: write
+
+jobs:
+  auto-merge:
+    runs-on: ubuntu-latest
+    if: github.actor == 'dependabot[bot]'
+    steps:
+      - name: Fetch Dependabot metadata
+        id: metadata
+        uses: dependabot/fetch-metadata@v2
+        with:
+          github-token: "${{ secrets.GITHUB_TOKEN }}"
+
+      - name: Auto-merge minor and patch updates
+        if: >-
+          steps.metadata.outputs.update-type == 'version-update:semver-minor' ||
+          steps.metadata.outputs.update-type == 'version-update:semver-patch'
+        run: gh pr merge --auto --squash "$PR_URL"
+        env:
+          PR_URL: ${{ github.event.pull_request.html_url }}
+          GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}

lmxlab-0.2.0/.github/workflows/docs.yml

@@ -0,0 +1,30 @@
+name: Docs
+
+on:
+  push:
+    branches: [main]
+
+permissions:
+  contents: read
+  pages: write
+  id-token: write
+
+jobs:
+  deploy:
+    runs-on: ubuntu-latest
+    environment:
+      name: github-pages
+      url: ${{ steps.deployment.outputs.page_url }}
+    steps:
+      - uses: actions/checkout@v6
+      - uses: astral-sh/setup-uv@v7
+        with:
+          enable-cache: true
+          cache-dependency-glob: "uv.lock"
+      - run: uv sync --locked --extra dev
+      - run: uv run mkdocs build
+      - uses: actions/upload-pages-artifact@v4
+        with:
+          path: site
+      - id: deployment
+        uses: actions/deploy-pages@v4

lmxlab-0.2.0/.github/workflows/publish.yml

@@ -0,0 +1,25 @@
+name: Publish to PyPI
+
+on:
+  release:
+    types: [published]
+
+permissions:
+  id-token: write
+
+jobs:
+  publish:
+    runs-on: ubuntu-latest
+    environment: pypi
+    steps:
+      - uses: actions/checkout@v6
+
+      - uses: astral-sh/setup-uv@v7
+        with:
+          enable-cache: true
+          cache-dependency-glob: "uv.lock"
+
+      - run: uv build
+
+      - name: Publish to PyPI
+        uses: pypa/gh-action-pypi-publish@release/v1

lmxlab-0.2.0/.gitignore

@@ -0,0 +1,42 @@
+# Python
+__pycache__/
+*.pyc
+*.pyo
+dist/
+*.egg-info/
+
+# Tool caches
+.ruff_cache/
+.mypy_cache/
+.pytest_cache/
+
+# Virtual environment
+.venv/
+
+# Claude Code (local-only)
+.claude/
+CLAUDE.md
+
+# Secrets
+.env
+.env.*
+credentials.json
+
+# Data (downloaded by recipes)
+/data/
+
+# Experiment results
+results.jsonl
+
+# Agent team output (temporary)
+.team-output/
+
+# Checkpoints
+checkpoints/
+*.safetensors
+
+# Build artifacts
+site/
+
+# OS
+.DS_Store

lmxlab-0.2.0/.pre-commit-config.yaml

@@ -0,0 +1,29 @@
+repos:
+  - repo: https://github.com/pre-commit/pre-commit-hooks
+    rev: v5.0.0
+    hooks:
+      - id: trailing-whitespace
+      - id: end-of-file-fixer
+      - id: check-yaml
+      - id: check-toml
+      - id: check-added-large-files
+        args: ["--maxkb=500"]
+
+  - repo: https://github.com/astral-sh/ruff-pre-commit
+    rev: v0.11.3
+    hooks:
+      - id: ruff
+        args: [--fix]
+      - id: ruff-format
+
+  - repo: https://github.com/astral-sh/uv-pre-commit
+    rev: 0.7.12
+    hooks:
+      - id: uv-lock
+
+  - repo: https://github.com/compilerla/conventional-pre-commit
+    rev: v4.0.0
+    hooks:
+      - id: conventional-pre-commit
+        stages: [commit-msg]
+        args: [feat, fix, refactor, test, docs, chore, ci, build, perf]

lmxlab-0.2.0/CHANGELOG.md

@@ -0,0 +1,92 @@
+# Changelog
+
+All notable changes to lmxlab will be documented in this file.
+
+The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/),
+and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
+
+## [Unreleased]
+
+## [0.2.0] - 2026-03-14
+
+### Added
+
+- **16 new architecture config factories**: DeepSeek V3 (MLA + MoE),
+  Nemotron (hybrid Mamba-Transformer MoE), Llama 4 Scout/Maverick
+  (iRoPE + chunked attention + MoE), Mistral Small (sliding window),
+  OLMo 2 (QK-norm), GPT-OSS (QK-norm), Grok (SharedExpertMoE),
+  Kimi K2.5 (DeltaNet + MoE), Qwen-Next (gated attention),
+  SmolLM3 (iRoPE), Qwen 3 MoE, Falcon H1 (hybrid Mamba-2),
+  Jamba (Mamba-2 + MoE), Bamba (hybrid Mamba-2), GLM-4.5 (MLA NoPE)
+- **Mamba-2 SSD**: structured state-space sequence mixer with chunked
+  parallel scan and recurrent inference paths
+- **Mamba-3**: trapezoidal discretization, BCNorm, complex A
+  (data-dependent RoPE on B/C)
+- **QK-norm**: per-head RMSNorm on Q and K projections (OLMo 2 style)
+- **GatedGQA**: sigmoid output gating on attention
+  (arXiv:2505.06708)
+- **ChunkedGQA**: fixed-size local attention with per-chunk RoPE
+  (Llama 4 iRoPE pattern)
+- **LatentMoE**: down-project before routing for many-expert MoE
+  (arXiv:2601.18089)
+- **SharedExpertMoE**: shared expert alongside routed experts
+  (DeepSeek V3 style)
+- **ReluSquaredFFN**: squared ReLU activation (Primer / Nemotron)
+- **muP parameterization**: width-independent hyperparameter transfer
+- **Dropout support**: configurable dropout in attention and FFN
+- **SparseGQA (DSA)**: DeepSeek Sparse Attention with compressed tokens,
+  selected tokens, and sliding window (arXiv:2512.02556)
+- **GRPOTrainer**: full GRPO training loop with group sampling, reward
+  scoring, and clipped surrogate objective (arXiv:2501.12948)
+- **Beam search**: standard beam search with optional custom scoring
+- **RewardModel**: language model + scalar head for reward scoring
+
+## [0.1.0] - 2026-03-11
+
+Initial release.
+
+### Added
+
+- **8 architecture config factories**: GPT, LLaMA, Gemma, Qwen, Mixtral (MoE),
+  DeepSeek V2 (MLA), Gemma 3 (sliding window), Qwen 3.5 (hybrid DeltaNet)
+- **Core components**: MHA, GQA, MLA, GatedDeltaNet, SlidingWindowGQA,
+  StandardFFN, GatedFFN, MoEFFN, SharedExpertMoEFFN, RMSNorm, LayerNorm,
+  RoPE, ALiBi, sinusoidal positional encoding
+- **ConfigurableBlock** with typed Registry pattern for component resolution
+- **LanguageModel** base class with tied/untied embeddings and KV cache
+- **Compiled training** with `mx.compile`, `nn.value_and_grad`, gradient
+  clipping, and cosine/linear/warmup learning rate schedules
+- **Optimizers**: AdamW, Lion, Adafactor, SGD with momentum
+- **Advanced training**: DPO, GRPO, multi-token prediction, curriculum
+  learning, knowledge distillation
+- **LoRA and QLoRA**: parameter-efficient fine-tuning with optional 4-bit
+  quantization
+- **Post-training quantization**: 4-bit and 8-bit via MLX native quantization,
+  with dequantization support
+- **Inference**: autoregressive generation with KV cache, streaming generation,
+  top-k/top-p/temperature sampling, repetition penalty, stop tokens
+- **Advanced inference**: best-of-N sampling, majority vote, speculative
+  decoding
+- **HuggingFace integration**: load pretrained weights (`load_from_hf`),
+  tokenizer wrapper (`HFTokenizer`), streaming dataset (`HFDataset`)
+- **Data pipeline**: CharTokenizer, TiktokenTokenizer, TextDataset,
+  TokenDataset, batch iterator
+- **Evaluation**: perplexity, bits-per-byte, pass@k for code generation
+- **Experiment framework**: ExperimentRunner with time budgets, ExperimentLog,
+  grid/random hyperparameter sweeps, statistical analysis (confidence
+  intervals, Cohen's d, experiment comparison)
+- **MLX profiling**: benchmark_fn, memory_estimate, profile_forward,
+  profile_generation, count_parameters_by_module
+- **CLI**: `lmxlab list`, `lmxlab info`, `lmxlab count`
+- **Callbacks**: MetricsLogger, EarlyStopping, ThroughputMonitor
+- **Checkpointing**: save/load via safetensors with JSON metadata
+- **31 recipe scripts** covering training, fine-tuning, DPO, GRPO, MTP,
+  distillation, curriculum learning, architecture comparison, ablation
+  studies, quantization, speculative decoding, and more
+- **Documentation site** with MkDocs Material: architecture guides, MLX
+  idioms, model comparison, API reference, recipes index, devlog
+- **PyPI publish workflow** via trusted publishing (OIDC)
+
+[Unreleased]: https://github.com/michaelellis003/lmxlab/compare/v0.2.0...HEAD
+[0.2.0]: https://github.com/michaelellis003/lmxlab/compare/v0.1.0...v0.2.0
+[0.1.0]: https://github.com/michaelellis003/lmxlab/releases/tag/v0.1.0

lmxlab-0.2.0/CONTRIBUTING.md

@@ -0,0 +1,153 @@
+# Contributing to lmxlab
+
+Thanks for your interest in contributing! lmxlab is a research platform,
+so clarity and rapid iteration are valued over production optimization.
+
+## Setup
+
+```bash
+git clone https://github.com/michaelellis003/lmxlab.git
+cd lmxlab
+uv sync --extra dev
+uv run pre-commit install
+uv run pre-commit install --hook-type commit-msg
+```
+
+The pre-commit hooks will automatically:
+- Run ruff lint and format checks on staged files
+- Verify `uv.lock` stays in sync with `pyproject.toml`
+- Enforce conventional commit message format
+
+## Development workflow
+
+1. Create a feature branch from `main`:
+   ```bash
+   git checkout -b feat/my-feature
+   ```
+
+2. Write tests first (TDD). Tests go in `tests/`:
+   ```bash
+   uv run pytest tests/test_my_module.py -v
+   ```
+
+3. Implement the feature in `src/lmxlab/`.
+
+4. Verify everything passes locally before pushing:
+   ```bash
+   uv run pytest                                    # All tests
+   uv run ruff check src/ tests/ recipes/            # Lint
+   uv run ruff format --check src/ tests/ recipes/   # Formatting
+   uv run mkdocs build --strict                      # Docs build
+   ```
+
+5. Open a PR against `main`. CI must pass before merging.
+
+## CI pipeline
+
+Every PR runs three jobs:
+
+- **lint** (ubuntu): ruff check + format on `src/`, `tests/`, `recipes/`
+- **docs** (ubuntu): `mkdocs build --strict` catches broken links/refs
+- **test** (macos-14): pytest on Apple Silicon (MLX requires M-series)
+
+All three must pass before merging. Do not bypass CI with `--admin`.
+
+## Keeping `uv.lock` in sync
+
+If you change `pyproject.toml` (add/remove/update dependencies), you must
+regenerate the lockfile:
+
+```bash
+uv lock
+```
+
+The `uv-lock` pre-commit hook catches this automatically. If CI fails with
+"lockfile needs to be updated", run `uv lock` and commit the result.
+
+## Branch naming
+
+- `feat/` — new features
+- `fix/` — bug fixes
+- `docs/` — documentation changes
+- `refactor/` — code restructuring
+- `test/` — test additions
+
+## Commit messages
+
+Follow the `type: description` format (enforced by pre-commit hook):
+
+```
+feat: add sliding window attention
+fix: correct RoPE dimension calculation
+docs: expand installation guide
+test: add CLI command tests
+refactor: simplify MoE routing logic
+```
+
+## Code style
+
+- **Line length:** 79 characters
+- **Quotes:** double quotes (enforced by ruff)
+- **Docstrings:** Google style
+- **Type annotations:** required on all public functions
+- **Imports:** sorted by ruff (stdlib, third-party, local)
+
+Ruff handles formatting and linting:
+```bash
+ruff check --fix src/ tests/ recipes/   # Auto-fix lint issues
+ruff format src/ tests/ recipes/         # Auto-format
+```
+
+## Testing
+
+- Use **behavioral tests** for ML code:
+  - Shape tests: output dimensions are correct
+  - Invariance tests: same input + seed = same output
+  - Directional tests: loss decreases after training
+  - Minimum functionality: no NaN, no Inf
+- Keep tests fast: use tiny model configs (`gpt_tiny()`, `llama_tiny()`)
+- Mark slow tests with `@pytest.mark.slow`
+
+## Architecture guidelines
+
+- **Config factories, not subclasses.** New architectures should be
+  config factory functions, not new model classes.
+- **Registry pattern.** New attention/FFN/norm types should register
+  themselves in the appropriate registry.
+- **Simplicity bias.** When two approaches achieve similar results,
+  prefer the simpler one.
+- **Clarity.** Comments should explain *why*, not just *what*.
+
+## Citations and attribution
+
+Every new building block (attention, FFN, position encoding, SSM,
+etc.) must cite its source paper with an arXiv ID in the module
+docstring:
+
+```python
+"""My new attention variant.
+
+Reference: Author (Year, arXiv:XXXX.XXXXX)
+"""
+```
+
+Code adapted from reference implementations must note the source:
+
+```python
+# Cross-references:
+# - org/repo filename.py (canonical implementation)
+# - HuggingFace transformers modeling_xxx.py
+```
+
+Use `Reference:` for the originating paper and `Cross-references:`
+for implementation sources consulted during development.
+
+## Adding a new architecture
+
+1. Create `src/lmxlab/models/myarch.py` with a config factory function
+2. Register any new components in the appropriate registry
+3. Add a `myarch_tiny()` config for tests
+4. Add tests in `tests/test_architectures.py`
+5. Update `src/lmxlab/models/__init__.py` exports
+6. Add to CLI in `src/lmxlab/cli.py`
+7. Document in `docs/models/index.md`

lmxlab-0.2.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 Michael Ellis
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.