feat: StateDict implementation for Conv2D

[bennibbelink]

• Implements StateDict and LoadStateDict on Conv2D, following the same pattern as Linear
• Adds state dict support to both MNIST example models and an --out flag to save trained models to file

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

[kolkov]

Thanks for the PR, Ben — the core implementation in conv2d.go is clean and matches Linear's pattern precisely. A few things to address before merge:

1. Silent assertion in TestConv2D_LoadStateDict_RoundTrip (must fix)

tolerance.AssertAllApproxEqual returns an error that describes the first mismatch. The current code ignores it, so the round-trip test never actually verifies anything — it will always pass even if the copy failed:

// Current — never actually checks
tolerance.AssertAllApproxEqual(origWeight, loadedWeight, tol)

// Fix
if err := tolerance.AssertAllApproxEqual(origWeight, loadedWeight, tol); err != nil {
    t.Errorf("weight round-trip mismatch: %v", err)
}

Same issue for the bias comparison below it.

2. Coverage gaps (nice to have, but the dtype checks are worth testing)

The Codecov report shows conv2d.go at 66.7% on this patch — the uncovered paths are the dtype validation checks and the "missing bias" error path. Could you add:

• TestConv2D_LoadStateDict_MissingBias (layer with bias, dict without)
• TestConv2D_LoadStateDict_WrongBiasShape
• TestConv2D_LoadStateDict_WrongDtype (for at least weight dtype)
• TestConv2D_LoadStateDict_RoundTrip_NoBias

3. Minor: printf format inconsistency between examples

mnist/main.go uses " ✓ Writing model to %s\n" but mnist-cnn/main.go uses "Writing model to %s\n" (no prefix). The CNN output block uses the same ✓ prefix style for all other lines (153-159) — should match.

4. Minor: doc comments on example model methods

The existing methods (Forward, Parameters, String) all have doc comments. The new StateDict() and LoadStateDict() don't. Not a lint failure (godot/revive are disabled for examples), but worth adding for consistency within the same file.

Everything else looks great. The import migration to the public nn/tensor packages in model.go is the right call, and strings.CutPrefix for key routing is idiomatic.

[bennibbelink]

Thanks @kolkov, I believe I addressed all your feedback. Coverage for the changes now at 100%

Feel free to squash on merge

[kolkov]

All four items addressed — silent assertion fixed, full error-path coverage added (missing bias, wrong shapes, wrong dtypes for both weight and bias, no-bias round-trip), printf prefix aligned, doc comments in place. Coverage at 100%. Clean work, thanks Ben.