Merge pull request #8 from AbdelStark/claude/quality-audit-hardening-2XHlk

test(core): quality audit — 18 new tests, fix stale docs, harden .gitignore

Author: AbdelStark

.gitignore

@@ -17,9 +17,22 @@ target
 # Contains mutation testing data
 **/mutants.out*/
 
-# RustRover
-#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
-#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
-#  and can be added to the global gitignore or merged into this file.  For a more nuclear
-#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
-#.idea/
+# IDE / Editor files
+.idea/
+.vscode/
+*.swp
+*.swo
+*~
+
+# OS files
+.DS_Store
+Thumbs.db
+
+# Model checkpoint artifacts
+*.mpk
+*.bin
+*.onnx
+*.safetensors
+
+# Criterion benchmark output
+target/criterion/

AGENTS.md

@@ -29,7 +29,7 @@ src/
 ├── serialization.rs    # Model weight save/load (NamedMpk, binary, compact formats)
 ├── two_phase.rs        # Two-phase inference primitives (phase1_batched, online_softmax_merge)
 ├── attention.rs        # Multi-head self-attention
-├── feed_forward.rs     # SwiGLU-style MLP
+├── feed_forward.rs     # Two-layer MLP with GELU activation
 └── utils.rs            # Causal mask generation helpers
 
 tests/
@@ -55,7 +55,7 @@ fixtures/                # Reference outputs from PyTorch
 
 ```bash
 cargo build                        # Build the project
-cargo test --all-features          # Run all 66 tests
+cargo test --all-features          # Run all 84 tests
 cargo test test_name               # Run specific test
 cargo clippy -- -D warnings        # Lint (warnings = errors)
 cargo fmt                          # Format code

README.md

@@ -83,13 +83,13 @@ cargo bench                        # Benchmarks
 
 ## Current Status
 
-**Alpha** (v0.1.0). Core algorithm implemented and tested with 57 passing tests (unit, differential, property-based, integration). Built on burn 0.20. Suitable for research and experimentation. Not yet suitable for production training at scale.
+**Alpha** (v0.2.0). Core algorithm implemented and tested with 84 passing tests (unit, differential, property-based, integration, doctest). Built on burn 0.20. Serialization (NamedMpk, binary, compact/half-precision) and two-phase inference integrated. Suitable for research and experimentation. Not yet suitable for production training at scale.
 
 Known limitations:
-- No weight serialization/loading (safetensors support planned)
-- Two-phase inference optimization is implemented but not integrated into the main forward pass
+- No PyTorch checkpoint import (safetensors format)
 - NdArray backend only tested; GPU backends (wgpu, CUDA, Metal) untested
 - No distributed training support
+- No KV-cache for autoregressive generation
 
 See [ROADMAP.md](ROADMAP.md) for planned features and progress.
 

ROADMAP.md

@@ -1,8 +1,8 @@
 # attnres-rs Roadmap
 
-## Current Phase: Alpha (v0.1.0)
+## Current Phase: Alpha (v0.2.0)
 
-Core algorithm implemented and tested. Suitable for research and experimentation.
+Core algorithm, serialization, and two-phase inference implemented. 84 tests passing. Suitable for research and experimentation.
 
 ---
 
@@ -29,7 +29,7 @@ Core algorithm implemented and tested. Suitable for research and experimentation
 - [x] Config save/load (JSON via burn's Config trait)
 - [x] Integrate two-phase inference into main `forward_two_phase` method
 - [x] Layer accessor methods for two-phase inference components
-- [x] 66 tests passing (unit, differential, property-based, integration, doctest)
+- [x] 84 tests passing (unit, differential, property-based, integration, doctest)
 - [ ] Pre-trained weight loading from PyTorch checkpoints
 - [ ] Model export utilities
 

tests/property_tests.rs

@@ -102,4 +102,68 @@ proptest! {
         let output = op.forward(&blocks, &partial);
         prop_assert_eq!(output.dims(), [batch, seq_len, d_model]);
     }
+
+    /// Zero-init AttnRes with identical sources should return that source.
+    /// This is the identity property: if all inputs are the same, the output
+    /// must equal that input regardless of num_blocks.
+    #[test]
+    fn identical_sources_produce_identical_output(
+        num_blocks in 1_usize..5,
+        batch in 1_usize..3,
+        seq_len in 1_usize..5,
+    ) {
+        let d_model = 16;
+        let device = Default::default();
+        let config = AttnResConfig::new(d_model, 12, num_blocks);
+        let op = config.init_op::<TestBackend>(&device);
+
+        // All sources are the same tensor
+        let source = Tensor::random(
+            [batch, seq_len, d_model],
+            Distribution::Uniform(-1.0, 1.0),
+            &device,
+        );
+        let blocks: Vec<_> = (0..num_blocks).map(|_| source.clone()).collect();
+
+        let output = op.forward(&blocks, &source);
+        let diff: f32 = (output - source).abs().max().into_scalar();
+        prop_assert!(
+            diff < 1e-3,
+            "Identical sources should produce that source, diff={diff}"
+        );
+    }
+
+    /// AttnRes output should be finite (no NaN/Inf) for any reasonable input.
+    #[test]
+    fn output_is_always_finite(
+        num_blocks in 1_usize..4,
+        batch in 1_usize..3,
+        seq_len in 1_usize..5,
+    ) {
+        let d_model = 16;
+        let device = Default::default();
+        let config = AttnResConfig::new(d_model, 12, num_blocks);
+        let op = config.init_op::<TestBackend>(&device);
+
+        let blocks: Vec<_> = (0..num_blocks)
+            .map(|_| {
+                Tensor::random(
+                    [batch, seq_len, d_model],
+                    Distribution::Uniform(-10.0, 10.0),
+                    &device,
+                )
+            })
+            .collect();
+        let partial = Tensor::random(
+            [batch, seq_len, d_model],
+            Distribution::Uniform(-10.0, 10.0),
+            &device,
+        );
+
+        let output = op.forward(&blocks, &partial);
+        let has_nan: bool = output.clone().is_nan().any().into_scalar();
+        let has_inf: bool = output.is_inf().any().into_scalar();
+        prop_assert!(!has_nan, "Output contains NaN");
+        prop_assert!(!has_inf, "Output contains Inf");
+    }
 }