1 files changed, 71 insertions, 0 deletions
diff --git a/tests/test_cuda_scaling.py b/tests/test_cuda_scaling.py
new file mode 100644
index 0000000..5f0befc
--- /dev/null
+++ b/tests/test_cuda_scaling.py
@@ -0,0 +1,71 @@
+"""Tests for CUDA device handling, chunked evaluation, and scaling to 1000 and 100_000."""
+import torch
+
+from src.config import Config, parse_args
+from src.data import build_examples, get_splits, next_prime, sieve_primes
+from src.model_api import build_model
+from src.train import evaluate, resolve_device
+
+
+def test_device_resolution():
+    assert resolve_device("cpu") == torch.device("cpu")
+    auto_dev = resolve_device("auto")
+    expected = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    assert auto_dev == expected
+
+
+def test_config_device_args():
+    cfg = parse_args(["rnn", "0", "--device", "cpu", "--eval_batch_size", "256"])
+    assert cfg.device == "cpu"
+    assert cfg.eval_batch_size == 256
+
+
+def test_bisect_next_prime_large_range():
+    primes = sieve_primes(100_100)
+    # Check boundary primes
+    assert next_prime(2, primes) == 3
+    assert next_prime(97, primes) == 101
+    assert next_prime(997, primes) == 1009
+    assert next_prime(99_991, primes) == 100_003
+
+
+def test_chunked_evaluate_matches_single_batch():
+    """Verify that chunked evaluation with eval_batch_size=4 gives identical results to full batch."""
+    cfg = Config(model="rnn", eval_batch_size=4)
+    model = build_model(cfg)
+    train_in, val_in = get_splits(cfg)
+    val_ex = build_examples(val_in, cfg)
+
+    # Chunked
+    tok_chunked, em_chunked, per_chunked = evaluate(model, val_ex, cfg)
+
+    # Single full batch
+    cfg_full = Config(model="rnn", eval_batch_size=len(val_ex))
+    tok_full, em_full, per_full = evaluate(model, val_ex, cfg_full)
+
+    assert tok_chunked == tok_full
+    assert em_chunked == em_full
+    assert per_chunked == per_full
+
+
+def test_range_extension_1000():
+    cfg = Config(range_start=2, range_end=1000, model="transformer")
+    train_in, val_in = get_splits(cfg)
+    assert len(train_in) + len(val_in) == 999
+    train_ex = build_examples(train_in, cfg)
+    assert len(train_ex) == len(train_in)
+    model = build_model(cfg)
+    tok, em, per = evaluate(model, build_examples(val_in[:20], cfg), cfg)
+    assert 0.0 <= em <= 1.0
+
+
+def test_range_extension_100k_data_shapes():
+    """Test data generation and model forward at 100k range without OOM."""
+    cfg = Config(range_start=2, range_end=100_000, model="rnn")
+    # Take a small slice of inputs from 100k range
+    sample_inputs = [2, 42, 999, 50_000, 99_991]
+    examples = build_examples(sample_inputs, cfg)
+    assert len(examples) == 5
+    model = build_model(cfg)
+    tok, em, per = evaluate(model, examples, cfg)
+    assert 0.0 <= em <= 1.0