Add collaboration docs scaffold with module index and issue/todo split

docs/agents.md

@@ -0,0 +1,22 @@
+# Agents
+
+Guidelines for LLM agents working on this repository.
+
+## Scope
+
+- Keep changes minimal and task-focused.
+- Preserve existing package and test structure.
+- Prefer factual documentation updates over speculative wording.
+
+## Working Rules
+
+- Read relevant source and tests before editing code.
+- Run the existing test command(s) available in the environment before and after changes when possible.
+- If tooling is unavailable in the environment, record that limitation in your notes.
+- Avoid broad refactors unless required by the issue.
+
+## Documentation Rules
+
+- Keep `docs/index.md` aligned with the current documentation tree.
+- Keep `docs/sources/index.md` aligned with current modules in `mlable/`.
+- Store proposals in `ideas.md`, selected work in `todo.md`, and curated GitHub issue focus in `issues.md`.

docs/context.md

@@ -0,0 +1,21 @@
+# Context
+
+## Project
+
+- Name: `mlable`
+- Package focus: TensorFlow utilities for layers, model blocks, metrics, sampling, schedules, and shaping helpers.
+- Packaging: Poetry project (`pyproject.toml`) targeting Python `>=3.10,<4.0`.
+
+## Source Layout
+
+- `mlable/blocks`: Attention and convolution building blocks, plus normalization and patching layers.
+- `mlable/layers`: Embedding, transformer, and shape-manipulation layers.
+- `mlable/maths`: Tensor reduction and probability helpers.
+- `mlable/models`: Generic contrast model, VAE model, and diffusion model classes.
+- `mlable/shaping`: Axis operations and Hilbert-curve folding utilities.
+- Top-level modules: caching, data, masking, metrics, sampling, schedules, shapes, text, utils.
+
+## Test Layout
+
+- Tests are in `tests/`.
+- The test tree mirrors source packages (`blocks`, `layers`, `maths`, `models`, `shaping`) plus top-level module tests.

docs/decisions.md

@@ -0,0 +1,9 @@
+# Decisions
+
+Important design and process decisions.
+
+## 2026-03-13
+
+- Added a repository-local `docs/` knowledge base for agent collaboration.
+- Added `docs/sources/index.md` to document each module in `mlable/`.
+- Kept `issues.md` as a curated subset of GitHub issues, with `todo.md` used for concrete next actions.

docs/ideas.md

@@ -0,0 +1,7 @@
+# Ideas
+
+Research directions and speculative items.
+
+- Add architecture diagrams for diffusion and autoencoder model flows.
+- Add short API examples for `sampling`, `schedules`, and `metrics` modules.
+- Add contribution playbooks for common tasks (new layer, new metric, new test).

docs/index.md

@@ -0,0 +1,38 @@
+# Index
+
+Overview of the repository structure.
+
+## Docs
+
+| Path                      | Purpose                               |
+| ------------------------- | ------------------------------------- |
+| `-- docs`                 | Collaboration-focused project docs    |
+| `   -- agents.md`         | Guidelines for LLM agents             |
+| `   -- context.md`        | Overview of the project               |
+| `   -- decisions.md`      | Record of important design choices    |
+| `   -- ideas.md`          | Research directions and speculations  |
+| `   -- index.md`          | Structure of the repository           |
+| `   -- invariants.md`     | Hard constraints                      |
+| `   -- issues.md`         | Curated list of selected issues       |
+| `   -- primer.jj`         | Conversation primer (Jinja template)  |
+| `   -- references.md`     | External references                   |
+| `   -- todo.md`           | Concrete next tasks                   |
+| `   -- sources/`          | Source-module documentation           |
+| `      -- index.md`       | Index of package modules              |
+
+## Sources
+
+| Path                      | Purpose                               |
+| ------------------------- | ------------------------------------- |
+| `-- mlable`               | Root of the Python package            |
+| `   -- blocks`            | Reusable model blocks                 |
+| `   -- layers`            | Standalone Keras layers               |
+| `   -- maths`             | Math and probability helpers          |
+| `   -- models`            | Model classes                         |
+| `   -- shaping`           | Axis and spatial transforms           |
+
+## Tests
+
+| Path                      | Purpose                               |
+| ------------------------- | ------------------------------------- |
+| `-- tests`                | Unit tests mirroring package layout   |

docs/invariants.md

@@ -0,0 +1,8 @@
+# Invariants
+
+Hard constraints for this repository.
+
+- Source code lives under `mlable/`.
+- Unit tests live under `tests/` and should mirror source structure.
+- Public package metadata and dependencies are defined in `pyproject.toml`.
+- Documentation in this folder should remain factual and concise.

docs/issues.md

@@ -0,0 +1,17 @@
+# Issues
+
+Curated list of selected GitHub issues for current focus.
+
+## Purpose
+
+`issues.md` is a manually curated shortlist, not a replacement for GitHub Issues.
+It should track only the issues currently prioritized for active work.
+
+## Current Selection
+
+- _No curated issues recorded yet._
+
+## Maintenance
+
+- Keep entries sorted by priority, then issue number.
+- Remove entries once merged or no longer relevant.

docs/primer.jj

@@ -0,0 +1,21 @@
+# Collaboration Primer
+
+You are assisting with the `mlable` repository.
+
+## Read First
+
+1. `docs/context.md`
+2. `docs/invariants.md`
+3. `docs/index.md`
+4. `docs/sources/index.md`
+
+## Current Focus
+
+- Selected issues: {{ selected_issues | default("See docs/issues.md") }}
+- Next tasks: {{ next_tasks | default("See docs/todo.md") }}
+
+## Response Requirements
+
+- Keep responses factual and concise.
+- Prefer minimal, targeted changes.
+- Preserve existing package/test structure.

docs/references.md

@@ -0,0 +1,8 @@
+# References
+
+External references used by this project.
+
+- TensorFlow API docs: <https://www.tensorflow.org/api_docs>
+- Keras API docs: <https://keras.io/api/>
+- Poetry docs: <https://python-poetry.org/docs/>
+- RoPE paper: <https://arxiv.org/pdf/2104.09864>

docs/sources/index.md

@@ -0,0 +1,51 @@
+# Sources Index
+
+Module-level overview for `mlable/`.
+
+## Top-Level Modules
+
+| Module | Purpose |
+| ------ | ------- |
+| `mlable/caching.py` | Cache creation and update helpers used by iterative decoding and attention. |
+| `mlable/data.py` | Dataset writing/statistics and codepoint decomposition helpers. |
+| `mlable/masking.py` | Contrast masking utilities for tensor comparisons. |
+| `mlable/metrics.py` | Grouped accuracy metrics for binary and categorical predictions. |
+| `mlable/sampling.py` | Top-k/top-p filtering and stochastic sampling functions. |
+| `mlable/schedules.py` | Numeric schedules, including linear and cosine rate helpers. |
+| `mlable/shapes.py` | Shape/dimension transformation helpers. |
+| `mlable/text.py` | UTF/codepoint text utilities and labeling helpers. |
+| `mlable/utils.py` | Generic utilities (composition, chunking, rotation, numeric helpers). |
+
+## Subpackages
+
+### `mlable/blocks`
+
+- `attention/generic.py`: attention block implementations and helpers.
+- `attention/transformer.py`: transformer-oriented attention block composition.
+- `convolution/generic.py`: generic convolutional block patterns.
+- `convolution/resnet.py`: ResNet-style convolutional blocks.
+- `convolution/unet.py`: U-Net style convolutional blocks.
+- `normalization.py`: adaptive group normalization layer.
+- `shaping.py`: patch extraction and shaping layer utilities.
+
+### `mlable/layers`
+
+- `embedding.py`: embedding and rotary embedding helper layers/functions.
+- `shaping.py`: Keras layers for divide/merge/swap/move axis transforms.
+- `transformer.py`: feed-forward transformer layer components.
+
+### `mlable/maths`
+
+- `ops.py`: reduction and grouping operations.
+- `probs.py`: probability utility functions (including log-normal PDF).
+
+### `mlable/models`
+
+- `autoencoder.py`: variational autoencoder model implementation.
+- `diffusion.py`: diffusion model base implementation.
+- `generic.py`: generic model utilities including contrast-oriented models.
+
+### `mlable/shaping`
+
+- `axes.py`: low-level axis divide/merge/swap/move operations.
+- `hilbert.py`: Hilbert-curve fold/unfold and permutation utilities.

docs/todo.md

@@ -0,0 +1,7 @@
+# Todo
+
+Concrete next tasks for repository maintenance.
+
+- [ ] Add a short contributor workflow section to `docs/context.md`.
+- [ ] Expand `docs/sources/index.md` with links to examples once available.
+- [ ] Keep curated issue shortlist in `docs/issues.md` synchronized with current priorities.