add rules, workflows and skills

Author: leondgarse

.agent/rules.md

@@ -0,0 +1,33 @@
+# Project Rules for keras_cv_attention_models
+
+## Code Style
+- Line length: 160 characters (`black -l 160`)
+- `__init__.py` files are excluded from black formatting
+- Use single-line multi-assignment for related vars: `self.a, self.b = a, b`
+
+## Model Architecture Patterns
+- Model variant files (e.g., `vit5.py`) should be **thin wrappers** calling the base function (e.g., `Beit`)
+- Use `@register_model` decorator for every model variant function
+- Use `kwargs.pop("kwargs", None)` at the start of base model functions to handle nested kwargs
+- Model names in code use **snake_case** (e.g., `vit5_small_patch16`), class-style names for functions (e.g., `ViT5_Small_Patch16`)
+- Default `pretrained="imagenet"` (or appropriate tag) once weights are available
+
+## PyTorch Backend (CRITICAL)
+- Any layer storing tensors as attributes (not via `add_weight`) **must** call `self.register_buffer()` in `__init__` with `if hasattr(self, "register_buffer")` guard
+- Prefer NumPy operations in `build()` for one-time tensor computation, then convert with `functional.convert_to_tensor()`
+- Always verify predictions match between TF and PyTorch backends after changes
+
+## Documentation
+- Every model family needs: paper link, GitHub link, model table with Params/FLOPs/Input/Top1 Acc/Download
+- Document all parameters in `__class_tail_doc__` with model-specific defaults noted (e.g., "Default True for ViT-5, False for others")
+- Download links: `https://github.com/leondgarse/keras_cv_attention_models/releases/download/{sub_release}/{filename}.h5`
+
+## Testing
+- Every model needs both `test_ModelFamily_defination` and `test_ModelFamily_predict` in `tests/test_models.py`
+
+## Modifying Existing Code
+- **Preserve existing patterns**: When refactoring a layer, study the original implementation's design choices (e.g., `blocks_shape` for dimension-agnostic reshaping) before replacing them with new approaches. Reuse proven patterns rather than inventing alternatives.
+- **Use proper APIs, not workarounds**: For PyTorch backend, use `register_buffer()` — don't set internal flags like `use_layer_as_module` directly. Always prefer the idiomatic solution over hacking internals.
+- **Run existing tests first**: Before modifying a layer, run its existing test (`test_layers.py`) to understand the input/output contract. Don't change behavior without checking what the tests expect.
+- **`build()` should use numpy**: Computations in `build()` are one-time setup. Use `np.tile`, `np.sin`, `np.concatenate` etc., then convert the final result with `functional.convert_to_tensor()`. Don't use `functional.tile` or other framework ops for build-time constants.
+- **Consolidate related operations**: When a helper function (e.g., `_get_sin_cos`) always has the same post-processing (e.g., tiling by `num_heads`), fold that into the helper rather than leaving it in the caller.

.agent/skills/weight-conversion/SKILL.md

@@ -0,0 +1,54 @@
+---
+name: Weight Conversion
+description: Converting PyTorch model weights to Keras h5 format for keras_cv_attention_models
+---
+
+# Weight Conversion Skill
+
+Convert pretrained PyTorch weights to Keras h5 format.
+
+The core task is **aligning the weight name order** between torch and keras models. `download_and_load.keras_reload_from_torch_model` is a convenience helper that automates this, but direct manual conversion is also fine.
+
+## Pipeline Overview
+
+1. **Weight collection** (`state_dict_stack_by_layer`): Groups torch `state_dict` entries by layer name (splitting on `.`), filtering via `skip_weights` and `unstack_weights`.
+2. **Name alignment** (`align_layer_names_multi_stage`): Reorders keras layer names to match torch weight order.
+3. **Weight transfer** (`keras_reload_stacked_state_dict`): Applies standard transforms (Conv2D/Dense transpose, etc.) plus custom `additional_transfer` overrides, then saves.
+
+## Parameter Reference
+
+### Weight Collection
+| Parameter | Purpose |
+|-----------|---------|
+| `skip_weights` | Weight name suffixes to drop (e.g., `["num_batches_tracked", "relative_position_index"]`) |
+| `unstack_weights` | Weights kept as individual entries instead of grouped with their layer (e.g., `["cls_token", "pos_embed", "gamma_1"]`) |
+
+### Name Alignment (order matching)
+| Parameter | Purpose |
+|-----------|---------|
+| `tail_align_dict` | Reposition layers by tail name: `{tail_name: offset}`. Negative offset moves layer earlier. Can be scoped by stack: `{"stack3": {"attn_gamma": -6}}` |
+| `full_name_align_dict` | Reposition by exact name: value can be negative offset, absolute position, or another layer's name string |
+| `tail_split_position` | Where to split name into head/tail (default `2`). E.g., `1` → head=`stack1`, tail=`attn_gamma` |
+| `specific_match_func` | Function returning the complete ordered name list, bypassing all alignment logic. Use for complex cases where dicts can't express the mapping |
+
+### Weight Transfer
+| Parameter | Purpose |
+|-----------|---------|
+| `additional_transfer` | Custom transforms: `{LayerClass: lambda ww: [...]}` or `{"name_suffix": lambda ww: [...]}`. Applied after default Conv2D/Dense transposes |
+
+## Workflow
+
+1. Create keras model with `pretrained=None, classifier_activation=None`
+2. Run with `do_convert=False` first to inspect both name lists
+3. Compare printed torch/keras weight lists — find misalignments
+4. Configure alignment parameters to fix ordering
+5. Run with `do_convert=True` — it predicts with both models and prints results
+6. Verify top prediction matches (usually `Egyptian_cat` for the cat test image)
+7. `md5sum output.h5` → add hash to `PRETRAINED_DICT`
+8. Upload to GitHub releases. Notify user if cannot upload directly.
+
+## Troubleshooting
+
+- **Shape mismatch**: Dense/Conv transposes are automatic; check if combined QKV needs `unstack_weights`
+- **Name ordering wrong**: Use `do_convert=False` to see lists side-by-side; adjust offsets or use `specific_match_func` for full control
+- **Predictions don't match**: Check `rescale_mode` in `add_pre_post_process()`, `classifier_activation`, or intermediate layer outputs

.agent/workflows/port-new-model.md

@@ -0,0 +1,160 @@
+---
+description: How to port a new model to keras_cv_attention_models
+---
+
+# Porting a New Model
+
+This workflow covers adding a new model architecture (vision or LLM) to `keras_cv_attention_models`.
+
+## 1. Research the Source Model
+
+1. Read the paper and reference implementation (usually PyTorch from GitHub).
+2. Identify which **existing model family** it's closest to:
+   - **Vision**: ViT-5 → `Beit`, EVA02 → `Beit`, DINOv2 → `Beit`, etc.
+   - **LLM**: LLaMA3 → `LLaMA2`, etc.
+3. Note key architectural differences: new layers, attention mechanisms, normalization, positional encoding.
+4. Check if required building blocks already exist in `attention_layers/` (e.g., `RMSNorm`, `PositionalEncodingFourierRot`, `ClassToken`, `CausalMask`).
+
+## 2. Implement New Components
+
+1. If new layers are needed, add them to the appropriate file in `attention_layers/`.
+2. **Reuse existing layers** whenever possible — check `attention_layers/__init__.py` for available components.
+3. If modifying existing layers (e.g., adding parameters to `PositionalEncodingFourierRot`):
+   - Ensure backward compatibility — existing models using the layer must still work.
+   - Add new parameters with defaults matching the old behavior.
+   - Update `get_config()` to include new parameters.
+
+> [!IMPORTANT]
+> **PyTorch backend compatibility**: Any layer that stores tensors as plain attributes (not via `add_weight`) must call `self.register_buffer()` in `__init__` for PyTorch backend support. Without this, the layer becomes a no-op identity in the PyTorch computation graph. See the PyTorch backend workflow for details.
+
+## 3. Create the Model Variant File
+
+Models that share a base architecture should be **thin wrappers**:
+
+**Vision model** (see `beit/vit5.py`, `beit/eva02.py`):
+```python
+from keras_cv_attention_models.beit.beit import Beit
+from keras_cv_attention_models.models import register_model
+
+def ModelFamily(model_name="model_family", **kwargs):
+    kwargs.pop("kwargs", None)
+    return Beit(**locals(), **kwargs)
+
+@register_model
+def ModelFamily_Small(input_shape=(224, 224, 3), num_classes=1000, activation="gelu",
+                      classifier_activation="softmax", pretrained="imagenet", **kwargs):
+    embed_dim = 384
+    num_heads = 6
+    return ModelFamily(**locals(), model_name="model_family_small", **kwargs)
+```
+
+**LLM model** (see `llama2/llama2.py`):
+```python
+from keras_cv_attention_models.llama2.llama2 import LLaMA2
+from keras_cv_attention_models.models import register_model
+
+@register_model
+def LLaMA3_8B(max_block_size=8192, vocab_size=128256, include_top=True,
+              activation="swish", pretrained=None, **kwargs):
+    num_blocks = 32
+    embedding_size = 4096
+    num_heads = 32
+    num_kv_heads = 8
+    return LLaMA2(**locals(), model_name="llama3_8b", **kwargs)
+```
+
+Key rules:
+- Use `@register_model` decorator for each variant.
+- Set `pretrained` default appropriately once weights are available.
+- Use `kwargs.pop("kwargs", None)` to handle nested kwargs.
+- Model names use **snake_case** (e.g., `vit5_small_patch16`, `llama2_42m`).
+
+## 4. Register in `__init__.py`
+
+1. Import the new model and variant functions.
+2. Add a `__head_doc__` string with paper link and GitHub link.
+3. Add a `__tail_doc__` for common parameters:
+   - **Vision**: `input_shape`, `num_classes`, `classifier_activation`, `pretrained`, plus model-specific params.
+   - **LLM**: `vocab_size`, `max_block_size`, `include_top`, `dropout`, `activation`, `pretrained`.
+4. Add model architecture table (Params, FLOPs, and model-specific columns).
+
+## 5. Weight Conversion
+
+1. Add entries to `PRETRAINED_DICT` in the base model file:
+   ```python
+   # Vision: {model_name: {pretrained_tag: {resolution: md5}}}
+   "vit5_small_patch16": {"imagenet": {224: "md5_hash_here"}},
+   # LLM: {model_name: {pretrained_tag: md5}} (no resolution key)
+   "llama2_42m": {"tiny_stories": "md5_hash_here"},
+   ```
+2. Convert weights using the appropriate method:
+   - **Vision**: `download_and_load.keras_reload_from_torch_model()` or a model-specific helper.
+   - **LLM**: `convert_huggingface_weights_to_h5()` or direct state_dict loading.
+3. After conversion, verify:
+   ```python
+   # Vision
+   mm = kecam.beit.ViT5_Small_Patch16(pretrained="path/to/converted.h5")
+   pred = mm(mm.preprocess_input(kecam.test_images.cat()))
+   print(mm.decode_predictions(pred)[0][0])  # ('n02124075', 'Egyptian_cat', 0.84...)
+
+   # LLM
+   mm = kecam.llama2.LLaMA2_42M(pretrained="path/to/converted.h5")
+   print(mm.run_prediction("A long time ago,", top_k=1, max_new_tokens=5))
+   ```
+4. Compute the md5 hash: `md5sum converted_model.h5`
+5. Upload to GitHub releases under the appropriate sub-release tag. Notice user if cannot upload directly.
+
+## 6. Documentation
+
+1. **Model-specific `README.md`** (e.g., `beit/README.md`, `llama2/README.md`):
+   - Add paper reference in the Summary section.
+   - Add model table with Params, FLOPs, and model-specific metrics.
+   - Download links: `[filename.h5](https://github.com/leondgarse/keras_cv_attention_models/releases/download/{sub_release}/{filename}.h5)`
+2. **Root `README.md`**: Add entry to Table of Contents in alphabetical order.
+3. **`__init__.py`**: Ensure all parameters are documented with model-specific defaults noted.
+
+## 7. Testing
+
+Add tests in `tests/test_models.py`:
+
+**Vision model**:
+```python
+def test_ModelFamily_defination():
+    mm = keras_cv_attention_models.beit.ModelFamily_Small(pretrained=None)
+    assert isinstance(mm, models.Model)
+    mm = keras_cv_attention_models.beit.ModelFamily_Small(pretrained=None, num_classes=0)
+    assert isinstance(mm, models.Model)
+
+def test_ModelFamily_predict():
+    mm = keras_cv_attention_models.beit.ModelFamily_Small(pretrained="imagenet")
+    pred = mm(mm.preprocess_input(cat()))
+    out = mm.decode_predictions(pred)[0][0]
+    assert out[1] == "Egyptian_cat"
+```
+
+**LLM model**:
+```python
+def test_LLM_run_prediction():
+    mm = keras_cv_attention_models.llama2.LLaMA2_42M(pretrained="tiny_stories")
+    generated = mm.run_prediction("A long time ago,", top_k=1, max_new_tokens=5)
+    assert generated == " there was a little girl"
+```
+
+If new layers were added, add layer tests in `tests/test_layers.py`.
+
+## 8. Format and Verify
+
+// turbo-all
+
+1. Run black formatter:
+   ```sh
+   find ./* -name "*.py" | grep -v __init__ | xargs -I {} black -l 160 {}
+   ```
+2. Run TF test suite:
+   ```sh
+   CUDA_VISIBLE_DEVICES='-1' pytest -vv --durations=0 ./tests
+   ```
+3. Run PyTorch test suite:
+   ```sh
+   CUDA_VISIBLE_DEVICES='-1' KECAM_BACKEND='torch' pytest -vv --durations=0 ./tests/test_models.py
+   ```

.agent/workflows/pytorch-backend-compat.md

@@ -0,0 +1,109 @@
+---
+description: PyTorch backend compatibility rules and gotchas for keras_cv_attention_models
+---
+
+# PyTorch Backend Compatibility
+
+The `pytorch_backend/` provides a custom Keras-like API on top of PyTorch. It uses a graph-based execution model where layers build a `GraphNode` computation graph during model construction, then execute it in `Model.forward()`. **This has several critical gotchas.**
+
+## Critical Rule: Layers with Non-Parameter Tensors
+
+**Problem**: In the PyTorch backend's `Layer.forward()` (`pytorch_backend/layers.py`), when a `GraphNode` input is encountered, the node's `callable` is set based on `self.use_layer_as_module`:
+
+```python
+if self.use_layer_as_module:  # Layer has weights → callable = self (invokes call())
+    cur_node.callable = self
+else:  # No weights → callable = self.module (identity lambda!)
+    cur_node.callable = self.module  # ← self.module = lambda xx: xx
+```
+
+`use_layer_as_module` is set to `True` only when `add_weight()` is called. **If a layer stores tensors as plain attributes without `add_weight()`, it becomes a silent no-op (identity function) in the PyTorch computation graph.**
+
+**Solution**: Use `register_buffer()` in `__init__` for any tensors that are not trainable parameters:
+
+```python
+def __init__(self, **kwargs):
+    super().__init__(**kwargs)
+    if hasattr(self, "register_buffer"):  # PyTorch backend only
+        self.register_buffer("my_tensor", None, persistent=False)
+        # Initialize as None, assign actual value in build()
+```
+
+This ensures:
+1. The layer is registered as a PyTorch submodule → `use_layer_as_module = True`
+2. Tensors are automatically moved to the correct device with `.to(device)`
+3. The layer's `call()` method is actually invoked during forward pass
+
+> [!CAUTION]
+> This is a **silent** failure — the model will run without errors but produce wrong predictions. The only symptom is degraded accuracy.
+
+## Diagnosing PyTorch Backend Issues
+
+### Symptom: Prediction mismatch between TF and PyTorch backends
+
+1. **Compare intermediate outputs** layer by layer:
+   ```python
+   # Enable debug mode to see layer-by-layer execution
+   mm.set_debug(True)
+   ```
+
+2. **Check if a layer is identity** by inspecting node callables:
+   ```python
+   for node in mm.forward_pipeline:
+       if 'layer_name' in node.name:
+           print(type(node.callable).__name__)  # 'function' = identity lambda
+           test_input = torch.randn(1, 197, 384)
+           print(torch.allclose(test_input, node.callable(test_input)))  # True = no-op!
+   ```
+
+3. **Write a comparison script** that runs both backends on the same input and compares named layer outputs.
+
+### Common Sources of Mismatch
+
+| Issue | Cause | Fix |
+|-------|-------|-----|
+| Layer is no-op | No `add_weight()` call → identity callable | Use `register_buffer()` in `__init__` |
+| Device mismatch | `convert_to_tensor()` creates CPU tensors | Use `register_buffer()` for auto device placement |
+| Shape mismatch | `functional.reshape` with hardcoded dims | Use `functional.shape(x)` or `blocks_shape` pattern for dimension-agnostic code |
+| Split behavior | `math.ceil` in integer splits | Check `pytorch_backend/functional.py:split()` |
+
+## Dimension-Agnostic Layer Design
+
+Layers should support both 3D `[batch, tokens, channels]` and 4D `[batch, heads, tokens, channels]` inputs. Use the `blocks_shape` pattern:
+
+```python
+def build(self, input_shape):
+    num_tokens = input_shape[-2]
+    self.blocks_shape = [*input_shape[1:-2], num_tokens]  # Captures intermediate dims
+
+def _process(self, inputs):
+    # Use blocks_shape for reshape — preserves batch and any intermediate dims
+    x = functional.reshape(inputs, [-1, *self.blocks_shape, self.channels // 2, 2])
+    ...
+    return functional.reshape(result, (-1, *self.blocks_shape, self.channels))
+```
+
+## Testing Both Backends
+
+Always run tests on both backends:
+
+```sh
+# TF backend (default)
+CUDA_VISIBLE_DEVICES='-1' pytest -vv --durations=0 ./tests
+
+# PyTorch backend
+CUDA_VISIBLE_DEVICES='-1' KECAM_BACKEND='torch' pytest -vv --durations=0 ./tests/test_models.py
+```
+
+## NumPy vs functional Operations in build()
+
+Prefer **NumPy** operations in `build()` for tensors that are computed once and stored (e.g., positional encodings). This avoids backend-specific issues with `functional.tile()`, `functional.convert_to_tensor()`, etc.:
+
+```python
+# Good: compute in numpy, convert once
+pos_sin = np.tile(np.sin(grid), [1, 1, num_heads])
+self.pos_sin = functional.convert_to_tensor(pos_sin, dtype=self.compute_dtype)
+
+# Risky: functional.tile() may behave differently across backends
+self.pos_sin = functional.tile(functional.convert_to_tensor(np.sin(grid)), [1, 1, num_heads])
+```

.agent/workflows/run-tests.md

@@ -0,0 +1,38 @@
+---
+description: How to run tests and formatting for keras_cv_attention_models
+---
+
+# Running Tests and Formatting
+
+// turbo-all
+
+## 1. Format Code with Black
+
+```sh
+find ./* -name "*.py" | grep -v __init__ | xargs -I {} black -l 160 {}
+```
+
+Note: `__init__.py` files are excluded from formatting.
+
+## 2. Run TF Backend Tests
+
+```sh
+CUDA_VISIBLE_DEVICES='-1' pytest -vv --durations=0 ./tests
+```
+
+This runs all tests: `test_models.py`, `test_layers.py`, `test_models_tf.py`, `test_switch_to_deploy_tf.py`.
+
+## 3. Run PyTorch Backend Tests
+
+```sh
+CUDA_VISIBLE_DEVICES='-1' KECAM_BACKEND='torch' pytest -vv --durations=0 ./tests/test_models.py
+```
+
+Only `test_models.py` is run for PyTorch backend.
+
+## Notes
+
+- These tests are also executed in GitHub Actions (`.github/workflows/publish-to-test-pypi.yml`), so running locally is not compulsory but recommended before pushing.
+- TF tests take ~14 minutes on CPU, PyTorch tests take ~2.5 minutes.
+- Use `CUDA_VISIBLE_DEVICES='-1'` to force CPU execution for reproducible results.
+- To run a single test: `pytest tests/test_models.py::test_ViT5_predict -xvs`