JSTPRV-165 Share initializer tensors

[millioner]

Description

Related Issue

Type of Change

• Bug fix (non-breaking)
• New feature (non-breaking)
• Breaking change (fix/feature causing existing functionality to break)
• Refactor (non-functional changes)
• Documentation update

Checklist

• Code follows project patterns
• Tests added/updated (if applicable)
• Documentation updated (if applicable)
• Self-review of code
• All tests pass locally
• Linter passes locally

Deployment Notes

Additional Comments

Summary by CodeRabbit

• Performance Improvements

  • Reduced memory usage by sharing tensor data across model layers.
  • Faster and more efficient quantization via a shared cache to avoid redundant work.

• Chores

  • Updated serialization dependency feature set to support improved runtime behavior.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro

Run ID: d522ea4d-8690-4b41-8167-c3aa3f57d52e

📥 Commits

Reviewing files that changed from the base of the PR and between 0694a41 and 02acecd.

📒 Files selected for processing (3)

• rust/jstprove_onnx/Cargo.toml
• rust/jstprove_onnx/src/graph.rs
• rust/jstprove_onnx/src/quantizer.rs

✅ Files skipped from review due to trivial changes (1)

• rust/jstprove_onnx/Cargo.toml

---

Walkthrough

Refactors tensor storage to use reference-counted Arcs and adds a shared quantization cache. Updates: LayerNode.weights now holds Arc<TensorData>, initializers and constant tensors are Arc-wrapped once, quantizer reuses cached Arcs, and serde gains the rc feature.

Changes

Cohort / File(s)	Summary
Dependency Configuration rust/jstprove_onnx/Cargo.toml	Added rc feature to the serde dependency (features = ["derive", "rc"]).
Graph Data Structure Refactor rust/jstprove_onnx/src/graph.rs	Changed LayerNode.weights: HashMap<String, TensorData> → HashMap<String, Arc<TensorData>>. Initializers and Constant node outputs are wrapped once in Arc and cloned (Arc::clone) into layers instead of deep-cloning tensors.
Quantizer Caching & Integration rust/jstprove_onnx/src/quantizer.rs	Added a weight_cache: HashMap<String, Arc<TensorData>> in quantize_model, passed to quantize_layer_weights. quantize_layer_weights reuses cached Arcs, skips quantization for already-integer tensors, constructs Arc-wrapped quantized tensors, and inserts folded BatchNorm tensors into layer.weights as Arc<TensorData>.

Sequence Diagram(s)

sequenceDiagram
    participant Parser as ONNX Parser
    participant Graph as LayerGraph
    participant Quant as Quantizer
    participant Cache as WeightCache

    Parser->>Graph: parse model.initializers & nodes
    Note right of Graph: wrap initializers/constants once\ninto Arc<TensorData>
    Graph->>Cache: populate cache entries for initializers (Arc)
    Graph->>Quant: provide layers with inputs referencing Arc weights
    Quant->>Cache: lookup weight by name
    alt cache hit
        Cache-->>Quant: Arc<TensorData> (clone)
        Quant->>Graph: insert Arc into layer.weights
    else cache miss & needs quantization
        Quant->>Quant: compute quantized int_data -> new TensorData
        Quant->>Cache: insert Arc(new TensorData)
        Cache-->>Quant: Arc(new TensorData)
        Quant->>Graph: insert Arc into layer.weights
    end

Loading

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~20 minutes

Possibly related PRs

• Extract ONNX ops into standalone jstprove_onnx crate #139: Modifies the same crate and modules (Cargo.toml, graph.rs, quantizer.rs); likely related introduction or prior changes to jstprove_onnx.
• Implement BatchNormalization folding and circuit constraint #114: Adds/follows BatchNormalization folding and quantization codepaths that this PR updates to store folded tensors as Arcs.

Poem

🐰 In tunnels of code I quietly hop,

One Arc for a tensor — no copies to drop,
Shared through the layers, a neat little map,
Caches snugly keep every quantized scrap,
My whiskers twitch — memory saved on the hop!

🚥 Pre-merge checks | ✅ 2 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (2 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title 'JSTPRV-165 Share initializer tensors' clearly and specifically describes the main change: refactoring to use Arc for shared initializer tensors across layers to reduce memory duplication.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

📝 Generate docstrings

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch shared_init_duplication

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}