fix(picklescan): safely parse bounded PyTorch tensor batches

AGENTS.md

@@ -18,7 +18,7 @@ This repo publishes **two PyPI packages with independent versions**:
 | `modelaudit`            | `./` (root)                       | `pyproject.toml` + `uv.lock`    | `CHANGELOG.md`                                |
 | `modelaudit-picklescan` | `packages/modelaudit-picklescan/` | `pyproject.toml` + `Cargo.toml` | `packages/modelaudit-picklescan/CHANGELOG.md` |
 
-Root `modelaudit` hard-requires `modelaudit-picklescan>=0.1.10,<0.2.0` — when the sibling crosses `0.2.0`, bump the constraint in the same PR or the next `modelaudit` release is uninstallable. Both packages are driven by a single `release-please` workflow (`.github/workflows/release-please.yml`) with components defined in `release-please-config.json` and current versions in `.release-please-manifest.json`. Full publishing details — trusted publishing, manual `workflow_dispatch` recovery (`root_version` / `picklescan_version`), and yank procedure — are in [`docs/agents/release-process.md`](docs/agents/release-process.md). For work inside the picklescan package, start from [`packages/modelaudit-picklescan/AGENTS.md`](packages/modelaudit-picklescan/AGENTS.md).
+Root `modelaudit` hard-requires `modelaudit-picklescan>=0.1.11,<0.2.0` — when the sibling crosses `0.2.0`, bump the constraint in the same PR or the next `modelaudit` release is uninstallable. Both packages are driven by a single `release-please` workflow (`.github/workflows/release-please.yml`) with components defined in `release-please-config.json` and current versions in `.release-please-manifest.json`. Full publishing details — trusted publishing, manual `workflow_dispatch` recovery (`root_version` / `picklescan_version`), and yank procedure — are in [`docs/agents/release-process.md`](docs/agents/release-process.md). For work inside the picklescan package, start from [`packages/modelaudit-picklescan/AGENTS.md`](packages/modelaudit-picklescan/AGENTS.md).
 
 ## Mission & Principles
 

CHANGELOG.md

@@ -7,6 +7,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Bug Fixes
+
+- Avoid false-positive PyTorch tensor reconstruction findings for bounded batched state dictionaries.
+
 ## [0.2.52](https://github.com/promptfoo/modelaudit/compare/v0.2.51...v0.2.52) (2026-07-22)
 
 ### Bug Fixes

docs/agents/release-process.md

@@ -11,7 +11,7 @@ Both packages are driven by a single [release-please](https://github.com/googlea
 
 The root release intentionally omits `package-name` and `component`. Release-please reads `project.name` from `pyproject.toml` for Python updates, while the empty branch component lets a root-only grouped Release PR match `release-please--branches--main` and keeps the existing `v{X.Y.Z}` tags. Restoring a non-empty `package-name` or `component` makes root-only releases look like a different component and silently skips publication after merge.
 
-The root `modelaudit` wheel declares a **hard dependency** on `modelaudit-picklescan>=0.1.10,<0.2.0` in `pyproject.toml`. When the sibling version crosses `0.2.0`, the constraint must be bumped in the same PR.
+The root `modelaudit` wheel declares a **hard dependency** on `modelaudit-picklescan>=0.1.11,<0.2.0` in `pyproject.toml`. When the sibling version crosses `0.2.0`, the constraint must be bumped in the same PR.
 
 ## Normal flow
 

packages/modelaudit-picklescan/AGENTS.md

@@ -4,7 +4,7 @@ Scoped agent guide for work inside `packages/modelaudit-picklescan/`. The root [
 
 ## What this package is
 
-`modelaudit-picklescan` is the Rust-backed pickle scanner that ships as an independent PyPI package. The root `modelaudit` wheel depends on it at runtime via a hard `modelaudit-picklescan>=0.1.10,<0.2.0` pin in the root `pyproject.toml`.
+`modelaudit-picklescan` is the Rust-backed pickle scanner that ships as an independent PyPI package. The root `modelaudit` wheel depends on it at runtime via a hard `modelaudit-picklescan>=0.1.11,<0.2.0` pin in the root `pyproject.toml`.
 
 - **Public API** — exported from `src/modelaudit_picklescan/__init__.py`: `PickleScanner`, `ScanOptions`, `scan_file`, `scan_bytes`, `scan_stream`, `shared_source_sensitive_caches`, `PickleReport`, `Finding`, `Notice`, `ScanError`, `Severity`, `ScanStatus`, `SafetyVerdict`, `CoverageSummary`. Treat these names as a stable surface.
 - **Rust engine** — `rust/src/` compiled to `modelaudit_picklescan._rust` via maturin + PyO3. Rust 1.83+, edition 2021.
@@ -74,7 +74,7 @@ Root-level validation (`uv run ruff check modelaudit/ packages/modelaudit-pickle
 - Release tag format: `modelaudit-picklescan-v{X.Y.Z}`.
 - Bumps are driven by Conventional Commits that **touch files inside `packages/modelaudit-picklescan/`**. Commits that only touch `modelaudit/` or the repo root do not bump this package.
 
-When this package reaches `0.2.0`, the root `pyproject.toml` `modelaudit-picklescan>=0.1.10,<0.2.0` constraint must be widened in the same PR, or the next `modelaudit` release will be uninstallable.
+When this package reaches `0.2.0`, the root `pyproject.toml` `modelaudit-picklescan>=0.1.11,<0.2.0` constraint must be widened in the same PR, or the next `modelaudit` release will be uninstallable.
 
 ## Publishing
 

packages/modelaudit-picklescan/CHANGELOG.md

@@ -7,6 +7,10 @@ and this package adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Bug Fixes
+
+- Validate bounded batched PyTorch state-dictionary entries without falsely flagging canonical tensor reconstruction.
+
 ## [0.1.10](https://github.com/promptfoo/modelaudit/compare/modelaudit-picklescan-v0.1.9...modelaudit-picklescan-v0.1.10) (2026-07-22)
 
 ### Bug Fixes

packages/modelaudit-picklescan/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "modelaudit-picklescan-rust"
-version = "0.1.10" # x-release-please-version
+version = "0.1.11" # x-release-please-version
 edition = "2021"
 rust-version = "1.83"
 description = "Native pickle security scanner engine for modelaudit-picklescan"

packages/modelaudit-picklescan/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "maturin"
 
 [project]
 name = "modelaudit-picklescan"
-version = "0.1.10" # x-release-please-version
+version = "0.1.11" # x-release-please-version
 description = "Standalone pickle security scanner extracted from ModelAudit"
 authors = [
     { name = "Ian Webster", email = "ian@promptfoo.dev" },
@@ -47,6 +47,7 @@ line-length = 120
 target-version = "py310"
 
 [tool.ruff.lint]
+select = ["E4", "E7", "E9", "F"]
 extend-select = [
     "I",
     "W",

packages/modelaudit-picklescan/src/modelaudit_picklescan/api.py

@@ -1932,14 +1932,14 @@ def apply_setitems_to_target(items: tuple[tuple[Any, Any], ...]) -> None:
         else:
             poison_stack_top()
 
-    def pop_marked_tuple() -> tuple[Any, ...] | None:
+    def pop_marked_tuple(*, max_width: int = _PYTORCH_STORAGE_TRUST_MAX_TUPLE_WIDTH) -> tuple[Any, ...] | None:
         items: list[Any] = []
         while stack:
             item = stack.pop()
             if item is marker:
                 return tuple(reversed(items))
             items.append(item)
-            if len(items) > _PYTORCH_STORAGE_TRUST_MAX_TUPLE_WIDTH:
+            if len(items) > max_width:
                 raise ValueError("PyTorch storage persistent ID tuple exceeded trust parser width")
         return None
 
@@ -2192,7 +2192,7 @@ def reduce_result(function: Any, args: Any, reduce_position: int) -> Any:
                 key = stack.pop()
                 apply_setitems_to_target(((key, value),))
             elif opcode_name == "SETITEMS":
-                setitem_items = pop_marked_tuple()
+                setitem_items = pop_marked_tuple(max_width=_PYTORCH_STORAGE_TRUST_MAX_STACK_DEPTH)
                 if setitem_items is None or len(setitem_items) % 2 != 0 or not stack:
                     clear_stack_after_malformed_provenance()
                     continue

packages/modelaudit-picklescan/tests/test_api.py

@@ -3432,6 +3432,92 @@ def test_scan_file_suppresses_rebuild_tensor_v2_for_real_ordered_state_dict(tmp_
     assert _torch_rebuild_tensor_v2_warning_dicts(report) == []
 
 
+def _write_large_batched_pytorch_state_dict(path: Path, *, malicious: bool = False) -> None:
+    entry_count = package_api._PYTORCH_STORAGE_TRUST_MAX_TUPLE_WIDTH // 2 + 1
+    entries: list[bytes] = []
+    for index in range(entry_count):
+        key_bytes = f"weight_{index}".encode("ascii")
+        key = b"X" + len(key_bytes).to_bytes(4, "little") + key_bytes
+        if malicious and index == entry_count - 1:
+            value = b"cos\nsystem\n(S'echo malicious-near-match'\ntR"
+        elif index == 0:
+            value = _pytorch_rebuild_tensor_v2_payload(key=str(index))[2:-1]
+        else:
+            value = _pytorch_rebuild_tensor_v2_reduce_expr(key=str(index))
+        entries.append(key + value)
+
+    payload = b"\x80\x04" + _global(b"collections", b"OrderedDict") + b")R(" + b"".join(entries) + b"u."
+    with zipfile.ZipFile(path, "w") as archive:
+        archive.writestr("archive/data.pkl", payload)
+        archive.writestr("archive/version", "3\n")
+        archive.writestr("archive/byteorder", "little")
+        for index in range(entry_count - int(malicious)):
+            archive.writestr(f"archive/data/{index}", b"\x00" * 24)
+
+
+def test_pytorch_storage_trust_parses_large_batched_state_dict_without_framework(tmp_path: Path) -> None:
+    archive_path = tmp_path / "large-batched-state.pt"
+    _write_large_batched_pytorch_state_dict(archive_path)
+    with zipfile.ZipFile(archive_path) as archive:
+        payload = archive.read("archive/data.pkl")
+
+    parsed = package_api._pytorch_storage_keys_from_pickle_bytes(payload)
+
+    assert parsed.parse_complete is True
+    assert len(parsed.referenced_keys) == package_api._PYTORCH_STORAGE_TRUST_MAX_TUPLE_WIDTH // 2 + 1
+    assert len(parsed.canonical_tensor_rebuild_invocations) == len(parsed.referenced_keys)
+
+
+def test_scan_file_suppresses_rebuild_tensor_v2_for_large_batched_state_dict(tmp_path: Path) -> None:
+    _require_torch_distribution()
+    archive_path = tmp_path / "large-batched-state.pt"
+    _write_large_batched_pytorch_state_dict(archive_path)
+
+    report = _scan_file_report_dict_subprocess(archive_path)
+
+    assert report["status"] == "complete"
+    assert report["verdict"] == "clean"
+    assert _torch_rebuild_tensor_v2_warning_dicts(report) == []
+    assert not any(finding["rule_code"] == "PERSISTENT_ID" for finding in report["findings"])
+
+
+def test_scan_file_preserves_malicious_call_in_large_batched_state_dict(tmp_path: Path) -> None:
+    archive_path = tmp_path / "large-batched-malicious-state.pt"
+    _write_large_batched_pytorch_state_dict(archive_path, malicious=True)
+
+    report = _scan_file_report_dict_subprocess(archive_path)
+
+    assert report["verdict"] == "malicious"
+    assert any(
+        finding["severity"] == "critical"
+        and finding["details"].get("module") in {"os", "posix", "nt"}
+        and finding["details"].get("name") == "system"
+        for finding in report["findings"]
+    )
+
+
+def test_pytorch_storage_trust_rejects_setitems_beyond_stack_limit() -> None:
+    entries = package_api._PYTORCH_STORAGE_TRUST_MAX_STACK_DEPTH // 2 + 1
+    oversized_payload = b"\x80\x04}(" + (b"K\x00K\x00" * entries) + b"u."
+
+    parsed = package_api._pytorch_storage_keys_from_pickle_bytes(oversized_payload)
+
+    assert parsed.parse_complete is False
+    assert parsed.referenced_keys == set()
+    assert parsed.canonical_tensor_rebuild_invocations == set()
+
+
+def test_pytorch_storage_trust_preserves_tuple_width_limit() -> None:
+    tuple_items = package_api._PYTORCH_STORAGE_TRUST_MAX_TUPLE_WIDTH + 1
+    oversized_payload = b"\x80\x04(" + (b"K\x00" * tuple_items) + b"t."
+
+    parsed = package_api._pytorch_storage_keys_from_pickle_bytes(oversized_payload)
+
+    assert parsed.parse_complete is False
+    assert parsed.referenced_keys == set()
+    assert parsed.canonical_tensor_rebuild_invocations == set()
+
+
 def test_scan_bytes_keeps_rebuild_tensor_v2_warning_for_raw_data_pickle(
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:

pyproject.toml

@@ -52,7 +52,7 @@ dependencies = [
     "posthog>=7.0.0",
     "protobuf>=5.29.0",
     "msgpack>=1.2.1,<2.0",
-    "modelaudit-picklescan>=0.1.10,<0.2.0",
+    "modelaudit-picklescan>=0.1.11,<0.2.0",
 ]
 
 [project.optional-dependencies]

tests/scanners/test_pytorch_zip_scanner.py

@@ -10003,6 +10003,21 @@ def test_pytorch_zip_scanner_suppresses_rebuild_tensor_v2_in_embedded_data_pkl(t
     assert _rebuild_tensor_v2_issue_dicts(report) == []
 
 
+def test_pytorch_zip_scanner_suppresses_rebuild_tensor_v2_in_large_batched_state_dict(tmp_path: Path) -> None:
+    _require_torch_distribution()
+    import torch
+
+    model_path = tmp_path / "large-batched-state.pt"
+    torch.save({f"weight_{index}": torch.zeros(1) for index in range(40)}, model_path)
+
+    report = _scan_pytorch_zip_report_dict_subprocess(model_path)
+
+    assert report["success"] is True
+    assert report["has_warnings"] is False
+    assert report.get("metadata", {}).get("pickle_verdict") == "clean"
+    assert _rebuild_tensor_v2_issue_dicts(report) == []
+
+
 def test_pytorch_zip_embedded_context_falls_back_with_older_picklescan_scan_stream(tmp_path: Path) -> None:
     _require_torch_distribution()
     from modelaudit_picklescan import PickleScanner as StandalonePickleScanner

tests/test_dependency_lock.py

@@ -14,7 +14,7 @@
 PICKLESCAN_PYPROJECT = ROOT_DIR / "packages" / "modelaudit-picklescan" / "pyproject.toml"
 PATCHED_GITPYTHON_FLOOR = (3, 1, 50)
 PINNED_MATURIN_BACKEND = "maturin===1.13.3"
-REQUIRED_PICKLESCAN_RELEASE = "modelaudit-picklescan>=0.1.10,<0.2.0"
+REQUIRED_PICKLESCAN_RELEASE = "modelaudit-picklescan>=0.1.11,<0.2.0"
 PATCHED_PY7ZR_REQUIREMENT = "py7zr>=1.1.3"
 PY7ZR_EXTRAS = ("sevenzip", "all-ci", "all")
 NUMPY_REQUIREMENTS = {