feat(streaming): AiDotNet-side weight streaming for PaLM-E 562B (addresses #1222)

[ooples]

Status: DRAFT — surface complete, but doesn't yet close the PaLM-E 562B OOM.

The streaming machinery in this PR works end-to-end for models whose weights fit in RAM at first forward (most production cases up through ~50B params). PaLM-E 562B specifically still OOMs because of an allocator-layer issue that needs a Tensors-side change. See "Remaining" below.

Summary

Lands the AiDotNet-side consumer for the weight-streaming infrastructure shipped in AiDotNet.Tensors 0.72.0 (PR #293 / Tensors). Addresses most of #1222; the PaLM-E 562B canary stays skipped until the missing Tensors-side allocator lands — issue stays open until then. (PR title intentionally uses "addresses #1222" rather than "closes #1222" so the issue isn't auto-closed at merge — closes review-comment #1271.yWer.)

Five tasks committed:

• [US-BF-019]: Replace NotImplementedException with ArgumentOutOfRangeException for unsupported kernel types #183 Auto-detect default streaming on NeuralNetworkBase ctor + post-first-forward retry (audit fix: previously the retry was placed before the forward and latched permanently before lazy weights materialized).
• [US-IF-007]: Verify acquisition functions and kernel types #184 Schedule-aware prefetch + MaterializeScope in PredictEager. Window W=2.
• fix(us-if-012): replace notimplementedexception with argumentoutofrangeexception for unsupported kernel types #185 ForwardForTraining reuses the streaming forward path; backward inherits LRU-warm tensors.
• fix(us-bf-006): fix LinearFeatureMapper training and usage issues #186 Public surface: IAiModelBuilder.ConfigureWeightStreaming (three-state Enabled + ThresholdParameters override that's actually wired now). AiModelResult.WeightStreamingReport populates from real WeightRegistry.GetStreamingReport() counters (audit fix: previously every counter was a 0 stub).
• fix(us-ci-001): replace intentional NotImplementedException with InvalidOperationException in transfer learning #187 11 regression tests (4 unit + 5 builder + 2 end-to-end). The end-to-end tests run an actual forward with streaming engaged — caught a stale-dll runtime failure during audit that surface-only tests had missed.

Hot path preserved: when streaming isn't engaged, PredictEager and ForwardForTraining take the simple foreach-and-forward fast paths bit-for-bit identical to pre-#1222.

Audit fixes (commit 190801e)

• Replaced the broken single-shot _streamingAutoDetectAttempted latch with _streamingAutoDetectFinalized + _streamingEngagedByAutoDetect + _firstForwardCompleted. Lazy networks now actually engage streaming on the post-first-forward retry.
• WeightStreamingReport.AutoDetected now distinguishes framework-engaged from user-forced (was always-on before).
• Real telemetry: counters populate from StreamingPoolReport (schema pinned via build-time probe). Removed the BytesWrittenToDisk/BytesReadFromDisk fields that the Tensors side doesn't expose.
• ConfigureWeightStreaming(config).ThresholdParameters actually drives the auto-detect comparison via NeuralNetworkBase.ApplyAutoDetectThresholdOverride. Was advertised but unwired.
• WeightRegistry.PrefetchAsync runtime invocation verified end-to-end (caught a MissingMethodException from a stale dll in test bin during audit).

Remaining: PaLM-E 562B requires a Tensors-side allocator

Root cause traced — MultiHeadAttentionLayer.OnFirstForward allocates weights as raw GC tensors:

_queryWeights  = new Tensor<T>([8192, 8192]);  // 537 MB
_keyWeights    = new Tensor<T>([8192, 8192]);  // 537 MB
_valueWeights  = new Tensor<T>([8192, 8192]);  // 537 MB
_outputWeights = new Tensor<T>([8192, 8192]);  // 537 MB
// 2.1 GB per MHA layer × 64 decoder layers = 134 GB before
// streaming has any chance to evict.

By the time RegisterTrainableParameter runs, the bytes are already on the GC heap. The streaming pool can DropStorageForStreaming (page to disk) but can't UN-allocate. Working set at peak still hits ~134 GB regardless of pool budget.

The fix needs:

1. Tensors v0.72.0 (separate PR): WeightRegistry.AllocateRegistered<T>(int[] shape) — atomically evicts LRU registered tensors to disk if needed to make headroom, allocates the new tensor, registers it with the pool.
2. AiDotNet (followup): route MultiHeadAttentionLayer.OnFirstForward, DenseLayer.OnFirstForward, EmbeddingLayer.OnFirstForward, Conv*.OnFirstForward through AllocateRegistered<T> instead of new Tensor<T>(...). Conditional on streaming being active so small models don't pay overhead.

Test plan

• dotnet build src/AiDotNet.csproj -f net10.0 — 0 errors
• dotnet test --filter "FullyQualifiedName~WeightStreaming" — 11/11 pass
• dotnet test --filter "FullyQualifiedName~LazyShape" — 45/45 pass (streaming branch is gated; small models take the fast path)
• CI green
• PaLM-E 562B canary — currently [Fact(Skip = "...")] in PaLMEProfilerTest. Still OOMs because of the eager-allocation issue above; un-skip after the Tensors v0.72.0 allocator + AiDotNet OnFirstForward routing land.

Bumps AiDotNet.Tensors 0.70.2 → 0.72.0

(Originally drafted against 0.71.0; rebased to 0.72.0 — closes review-comment #1271.yWeL on the Directory.Packages.props mismatch.)

🤖 Generated with Claude Code

Summary by CodeRabbit

• New Features

  • Per-build weight-streaming control and runtime telemetry surfaced on model results for large-model memory savings and diagnostics.
  • Broader lazy/streaming-aware layer initialization and added serialization hooks for more flexible parameter loading and reduced memory usage.
  • Safer large-model parameter handling to support very large parameter counts without overflow.

• Chores

  • Updated AiDotNet.Tensors central package to 0.72.0.

• Tests

  • Added unit and integration tests covering weight-streaming, auto-detection, and end-to-end streaming behavior.

[vercel]

The latest updates on your projects. Learn more about Vercel for GitHub.

Project	Deployment	Actions	Updated (UTC)
aidotnet_website	Error		May 6, 2026 3:48pm
aidotnet-playground-api	Error		May 6, 2026 3:48pm

[vercel]

Deployment failed with the following error:

The `vercel.json` schema validation failed with the following message: `ignoreCommand` should NOT be longer than 256 characters

Learn More: https://vercel.com/docs/concepts/projects/project-configuration

[coderabbitai]

Note

Reviews paused

It looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the reviews.auto_review.auto_pause_after_reviewed_commits setting.

Use the following commands to manage reviews:

• @coderabbitai resume to resume automatic reviews.
• @coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

• ▶️ Resume reviews
• 🔍 Trigger review

Walkthrough

Adds weight‑streaming feature (config, report DTO, builder API), streaming-aware lazy weight allocation and forward path, a ParameterCount helper for safe long→int flattening, widespread layer/adapter/model updates to use lazy allocation and helper sizing, test updates, and bumps AiDotNet.Tensors 0.70.2→0.72.0.

Changes

Weight streaming + Parameter-count / lazy-allocation (single cohesive change DAG)

Layer / File(s)	Summary
Data Shape / DTOs src/Deployment/Configuration/WeightStreamingConfig.cs, src/Deployment/Configuration/WeightStreamingReport.cs	Adds WeightStreamingConfig (nullable Enabled, optional ThresholdParameters) and WeightStreamingReport DTO with init-only telemetry fields.
Public API / Builder src/Interfaces/IAiModelBuilder.cs, src/AiModelBuilder.cs	Adds IWeightStreamingCapableBuilder<T,TInput,TOutput> and ConfigureWeightStreaming(WeightStreamingConfig? config = null) implemented on AiModelBuilder; builder stores per-build config, applies it early in BuildAsync, re-applies after AutoML/RL/model reassignment, and surfaces WeightStreamingReport on results.
Core runtime: auto-detect & routing src/NeuralNetworks/NeuralNetworkBase.cs	Introduces streaming auto-detect/state (TryAutoEnableWeightStreaming, env override AIDOTNET_STREAMING_THRESHOLD_PARAMS, disable/override flags), streaming-aware forward/training routing (PredictEagerStreaming, prefetch window, prefetch/materialize helpers), one-time post-first-forward retry, and reporting hooks.
Layer allocation surface src/NeuralNetworks/Layers/LayerBase.cs, src/NeuralNetworks/Layers/*	Adds AllocateLazyWeight(int[] shape, ...), UseStreamingAllocator and converts many layers to lazy/streaming-aware allocations (AllocateLazyWeight), lazy shape-resolution in SetParameters and added Serialize/Deserialize where needed.
Parameter-count flattening helper src/Helpers/ParameterCountHelper.cs	New internal helper ParameterCountHelper.ToFlatVectorSize(long) validating non-negative and ≤ int.MaxValue and returning int flat size (throws on overflow/negative).
Models / adapters / callers src/LoRA/Adapters/*, src/*Model*.cs, src/AutoML/*, src/Regression/*, src/Diffusion/*, src/PhysicsInformed/*	Widened ParameterCount arithmetic to long where needed; replaced (int)ParameterCount casts with ParameterCountHelper.ToFlatVectorSize(ParameterCount) across adapters, models and gradient/vector allocation sites. Minor adapter accessors added in a few places.
Result wiring src/Models/Options/AiModelResultOptions.cs, src/Models/Results/AiModelResult.cs	Adds WeightStreamingReport to result options and AiModelResult; preserves report across WithParameters, serialization/deserialization, and populates it for supervised/RL/AutoML/meta-learning paths.
Tests / Integration / E2E tests/AiDotNet.Tests/IntegrationTests/WeightStreaming/*, tests/AiDotNet.Tests/UnitTests/NeuralNetworks/WeightStreaming/*, tests/.../NeuralNetworks/*	Adds builder integration tests, auto-detect unit tests, end-to-end streaming tests, and updates many tests to resolve shapes before inspecting ParameterCount or parameters.
Validation / Safety src/NeuralNetworks/NeuralNetworkBase.cs, various SetParameters sites	Adds overflow guards, explicit errors when ParameterCount > int.MaxValue, idempotent auto-detect guards, and safe failure when ParameterCount throws during partial construction.
Package manifest Directory.Packages.props	Bumps AiDotNet.Tensors from 0.70.2 → 0.72.0.

Sequence Diagram(s)

sequenceDiagram
    participant User
    participant Builder as AiModelBuilder
    participant Model as NeuralNetworkBase
    participant Registry as WeightRegistry
    participant Executor as Trainer/Inferencer

    User->>Builder: ConfigureWeightStreaming(config)
    Builder->>Builder: store _weightStreamingConfig
    User->>Builder: BuildAsync()
    Builder->>Model: ApplyWeightStreamingConfig()
    alt config.Enabled == true
        Model->>Model: ConfigureWeightLifetime(GpuOffloadOptions)
    else config.Enabled == false
        Model->>Model: DisableAutoStreaming()
    else
        Model->>Model: defer auto-detect
    end

    Executor->>Model: EnsureArchitectureInitialized / Predict
    Model->>Model: TryAutoEnableWeightStreaming()
    Model->>Model: evaluate ParameterCount vs threshold
    alt engage streaming
        Model->>Model: set weight lifetime configured
    else do not engage
    end

    Executor->>Model: Predict/ForwardForTraining
    alt _weightLifetimeConfigured == true
        loop each layer i
            par Prefetch
                Model->>Registry: PrefetchLayerWeights(i+2)
            and Materialize+Forward
                Model->>Registry: BeginLayerMaterializeScope(i)
                Registry-->>Model: materialized tensors
                Model->>Model: Layer.Forward()
                Model->>Registry: Dispose scope -> release
            end
        end
    else
        Model->>Model: PredictEager (non-streaming)
    end

    Model->>Builder: request BuildWeightStreamingReport()
    Builder->>Registry: Get counters / snapshot
    Builder-->>User: AiModelResult including WeightStreamingReport

Loading

Estimated code review effort

🎯 5 (Critical) | ⏱️ ~150 minutes

Possibly related PRs

• ooples/AiDotNet#1225: Related—overlaps GPU offload / weight-streaming wiring and WeightRegistry interactions.
• ooples/AiDotNet#1219: Related—prior work widening ParameterCount to long and addressing lazy-init issues referenced by linked issue; high overlap.
• ooples/AiDotNet#1244: Related—introduces/uses ParameterCountHelper.ToFlatVectorSize and aligns parameter-count flattening across layers/adapters.

Suggested labels

feature

We pawned the heavies to disk, one layer at a time,
Prefetching futures so the forward run stays sane.
Reports in hand and thresholds checked with care —
Lazy weights, flat counts, and tests everywhere. 🎛️🚀

BLOCKING issues (production-readiness checklist)

• All stubs/placeholders/no-op implementations (e.g., NoOpDisposable or simplified no-op materialize scopes) must be audited and replaced or documented with production-safe semantics; treat any no-op that can mask materialize/evict failures as BLOCKING.
• WeightRegistry prefetch/materialize/metrics paths that swallow failures or return best-effort counters must provide deterministic error/observable paths (logging, surfaced errors, or fail-fast modes) for production; unresolved silent failures are BLOCKING.
• Test-only helpers or test-only extension methods exposed in non-test assemblies (e.g., ConfigureWeightLifetimeForTest) must be gated or removed; any test-only API reaching production code is BLOCKING.
• The ParameterCount helper throws on counts > int.MaxValue; callers must handle or document migration strategies. Any code paths that could surface this InvalidOperationException to users without mitigation must be addressed before merge. BLOCKING.
• Serialization changes (many layers add Serialize/Deserialize) require versioning and backward-compatibility round-trip tests across persisted formats; lack of migration/versioning is BLOCKING.
• Concurrency and thread-safety: ensure WeightRegistry, streaming allocator, prefetch/materialize, and counters are safe for concurrent Predict/Train calls; potential races during first-forward materialization are BLOCKING until proven safe.
• API compatibility: adding IWeightStreamingCapableBuilder and implementing it on AiModelBuilder preserves source/binary compatibility for existing implementers, but confirm third-party IAiModelBuilder implementers are unaffected and add compatibility tests. BLOCKING.
• Ensure the new env-var-based threshold parsing and any public default values are documented and tested for internationalization/format issues and invalid inputs (negative or non-numeric env values must fail clearly). BLOCKING.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch feat/1222-weight-streaming-aidotnet

[Copilot]

Pull request overview

Adds AiDotNet-side integration for AiDotNet.Tensors weight streaming to prevent OOM on extremely large models (e.g., PaLM‑E 562B), including auto-detection, streaming-aware forward paths, builder configuration, and regression tests.

Changes:

• Introduces auto-detect + opt-out hooks in NeuralNetworkBase and adds a streaming-aware forward path with prefetch + per-layer materialize scopes.
• Adds public builder surface ConfigureWeightStreaming(...) and threads a WeightStreamingReport onto AiModelResult.
• Adds unit/integration tests for the new streaming configuration surface and DTOs; bumps AiDotNet.Tensors to 0.71.0.

Reviewed changes

Copilot reviewed 10 out of 10 changed files in this pull request and generated 12 comments.

Show a summary per file

File	Description
tests/AiDotNet.Tests/UnitTests/NeuralNetworks/WeightStreaming/AutoDetectWeightStreamingTests.cs	Adds unit tests intended to pin auto-detect/opt-out behavior.
tests/AiDotNet.Tests/IntegrationTests/WeightStreaming/WeightStreamingBuilderIntegrationTests.cs	Adds integration tests for builder chaining and DTO shape.
src/NeuralNetworks/NeuralNetworkBase.cs	Adds eager + lazy auto-detect calls; adds streaming-aware forward path with prefetch/materialize scopes.
src/AiModelBuilder.cs	Adds builder configuration for weight streaming and plumbs a report onto build results.
src/Interfaces/IAiModelBuilder.cs	Exposes new fluent API ConfigureWeightStreaming(...) with documentation.
src/Models/Results/AiModelResult.cs	Adds WeightStreamingReport to results.
src/Models/Options/AiModelResultOptions.cs	Adds option passthrough for WeightStreamingReport.
src/Deployment/Configuration/WeightStreamingConfig.cs	Adds config DTO (Enabled tri-state + threshold).
src/Deployment/Configuration/WeightStreamingReport.cs	Adds report DTO for streaming telemetry.
Directory.Packages.props	Bumps AiDotNet.Tensors package version to 0.71.0.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[coderabbitai]

Actionable comments posted: 11

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@src/AiModelBuilder.cs`:
- Around line 5390-5435: BuildWeightStreamingReport currently returns hardcoded
zeroes for the streaming counters; replace the placeholders by calling the
model's WeightRegistry streaming report and mapping its counters into the
returned AiDotNet.Deployment.Configuration.WeightStreamingReport (e.g. call
nnBase.WeightRegistry?.GetStreamingReport(), null-check it, and assign
DiskReadCount, EvictionCount, PrefetchIssueCount, PrefetchHitCount,
PrefetchMissCount, BytesWrittenToDisk, BytesReadFromDisk from that report), and
wrap that call in a try/catch so any version mismatch or exception leaves the
fields at the safe zero fallback while preserving StreamingEnabled,
AutoDetected, ModelParameterCount and EffectiveThresholdParameters.
- Around line 5438-5471: ApplyWeightStreamingConfig currently accepts
configurations that cannot be enforced (per-instance ThresholdParameters and
turning off streaming when it may already be engaged); make it fail-fast: in
ApplyWeightStreamingConfig validate _weightStreamingConfig and throw a
NotSupportedException for unsupported/unenforceable intents instead of silently
no-op — specifically throw if _weightStreamingConfig.Enabled == null and
_weightStreamingConfig.ThresholdParameters is set, and throw if
_weightStreamingConfig.Enabled == false (since
DisableAutoStreaming/ConfigureWeightLifetime cannot guarantee forcing an
already-engaged stream off); keep only the supported path that calls
nnBase.ConfigureWeightLifetime(new GpuOffloadOptions()) when Enabled == true and
return early for null config.

In `@src/Interfaces/IAiModelBuilder.cs`:
- Around line 1238-1241: The example in IAiModelBuilder.cs is misleading because
calling ConfigureWeightStreaming() with no arguments passes null and preserves
auto-detection rather than forcing streaming; update the example to explicitly
opt in by calling ConfigureWeightStreaming(true) (or the API's explicit enable
parameter) after ConfigureModel(myLargeModel) so streaming is actually forced;
locate the ConfigureWeightStreaming method and the example block in
IAiModelBuilder.cs and change the parameterless invocation to the explicit true
argument (or the named enable parameter) and adjust the comment to state it
explicitly enables weight streaming.

In `@src/Models/Options/AiModelResultOptions.cs`:
- Around line 745-756: The XML docs for the property WeightStreamingReport on
AiModelResultOptions are missing the required <value> element and the
`<para><b>For Beginners:</b>` subsection inside `<remarks>` per the Options
golden pattern; update the documentation for the public property
WeightStreamingReport to include a concise `<value>` describing the property
type and return semantics (nullable WeightStreamingReport), and extend the
existing `<remarks>` to add a `<para><b>For Beginners:</b>` that plainly
explains what enabling weight streaming does, when this property will be
populated (e.g., only when streaming during a build), and how callers should use
AiModelResult.WeightStreamingReport forwarded via the options constructor.

In `@src/Models/Results/AiModelResult.cs`:
- Line 1351: The new field WeightStreamingReport is assigned from options but
not preserved across object reconstruction and cloning; update Deserialize(...)
to read and set WeightStreamingReport from the serialized data, update
WithParameters(...) to forward the existing WeightStreamingReport into the
returned instance, and update DeepCopy(...) to copy/clone WeightStreamingReport
(or pass it into the copy constructor/options) so the report is not lost after
deserialize/with/clone operations; locate usages of WeightStreamingReport, the
options constructor where it’s set, and the methods Deserialize, WithParameters,
and DeepCopy to make these changes.

In `@src/NeuralNetworks/NeuralNetworkBase.cs`:
- Around line 3491-3521: TryAutoEnableWeightStreaming incorrectly marks
auto-detect as attempted when ParameterCount is 0 (lazy models before weights
exist), preventing a later forward from rechecking; change the logic in
TryAutoEnableWeightStreaming so that after successfully reading ParameterCount
you treat a value of 0 as "indeterminate" and return without setting
_streamingAutoDetectAttempted (i.e., only set _streamingAutoDetectAttempted when
ParameterCount is non-zero and you either decide to leave eager or enable
streaming), keeping the existing threshold check against
s_streamingThresholdParams and preserving the explicit ConfigureWeightLifetime
entry point.
- Around line 2701-2717: ForwardForTraining() currently returns early from the
checkpointing branch and never honors _weightLifetimeConfigured, causing
streaming training to be bypassed; modify ForwardForTraining() so that when
checkpointing is active AND _weightLifetimeConfigured is true it either calls
PredictEagerStreaming(input) within the checkpointing branch (i.e., compose
checkpointing over the streaming/materialize path by routing the checkpoint
branch to use PredictEagerStreaming) or, if composition is not yet implemented,
throw a clear NotSupportedException indicating checkpointing+weight streaming is
not supported; update any checkpoint-related helper (the checkpoint branch code)
to forward to PredictEagerStreaming instead of the non-streaming forward so the
same MaterializeScope/LRU behavior is preserved.
- Around line 2938-2971: PredictEagerStreaming can materialize lazy layer
weights during Layers[i].Forward(current) but those newly-allocated tensors are
never re-registered, so prefetch/materialize later still miss them; after the
using (BeginLayerMaterializeScope(i)) block (or immediately after
Layers[i].Forward returns while still in the materialize scope) call the routine
that registers weight lifetimes for that layer (e.g., invoke the same logic as
ConfigureWeightLifetime/RegisterTrainableTensorsWithWeightRegistry for layer i,
or add a new Refresh/RegisterWeightEntriesForLayer(layerIndex) helper) so any
tensors with Length>0 allocated during Forward are recorded in the weight
registry before the loop proceeds and before the scope is disposed.

In
`@tests/AiDotNet.Tests/IntegrationTests/WeightStreaming/WeightStreamingBuilderIntegrationTests.cs`:
- Around line 44-54: The test
ConfigureWeightStreaming_AcceptsAllThreeEnabledStates currently only verifies
"no throw"; change it to assert the config was actually applied: after each call
to builder.ConfigureWeightStreaming(new WeightStreamingConfig { Enabled = ...
}), assert that the builder's stored WeightStreamingConfig.Enabled equals the
value passed (null/true/false). Use the public accessor if one exists (e.g., a
property/method exposing WeightStreamingConfig) or, if none, read the backing
field (e.g., _weightStreamingConfig) via reflection to compare the Enabled
value; keep references to AiModelBuilder, ConfigureWeightStreaming and
WeightStreamingConfig.Enabled.

In
`@tests/AiDotNet.Tests/UnitTests/NeuralNetworks/WeightStreaming/AutoDetectWeightStreamingTests.cs`:
- Around line 80-110: The test currently performs DisableAutoStreaming and
TryAutoEnableWeightStreaming but has no assertion; add a concrete assertion that
the network did not enable streaming after the opt-out: after calling
net.DisableAutoStreaming(); net.TryAutoEnableWeightStreaming(); assert that
net.WeightStreamingAutoDetected is false (or the public equivalent
property/method on FixedParamCountNetwork) so the test fails if auto-enable
overrides the explicit opt-out.
- Around line 113-125: The test currently only checks
WeightStreamingAutoDetected but not whether ParameterCount was re-read; update
the test to instrument ParameterCount reads and assert it was evaluated exactly
once: create or use an instrumented network (e.g., subclass
FixedParamCountNetwork or a test-double) that increments a read-counter inside
the ParameterCount property, call TryAutoEnableWeightStreaming multiple times as
before, and assert the counter equals 1 and WeightStreamingAutoDetected remains
unchanged; reference FixedParamCountNetwork.ParameterCount,
TryAutoEnableWeightStreaming(), and WeightStreamingAutoDetected when adding the
counter and final assertion.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: ASSERTIVE

Plan: Pro

Run ID: cffa2023-5251-4002-b62a-ce82926b569e

📥 Commits

Reviewing files that changed from the base of the PR and between 5f6579d and 044fc07.

📒 Files selected for processing (10)

• Directory.Packages.props
• src/AiModelBuilder.cs
• src/Deployment/Configuration/WeightStreamingConfig.cs
• src/Deployment/Configuration/WeightStreamingReport.cs
• src/Interfaces/IAiModelBuilder.cs
• src/Models/Options/AiModelResultOptions.cs
• src/Models/Results/AiModelResult.cs
• src/NeuralNetworks/NeuralNetworkBase.cs
• tests/AiDotNet.Tests/IntegrationTests/WeightStreaming/WeightStreamingBuilderIntegrationTests.cs
• tests/AiDotNet.Tests/UnitTests/NeuralNetworks/WeightStreaming/AutoDetectWeightStreamingTests.cs

[coderabbitai]

Actionable comments posted: 5

♻️ Duplicate comments (4)

src/NeuralNetworks/NeuralNetworkBase.cs (2)

2669-2713: ⚠️ Potential issue | 🟠 Major | ⚡ Quick win

Checkpointing still bypasses the streaming training path.

When the checkpoint branch at Lines 2669-2713 returns early, _weightLifetimeConfigured is never consulted, so enabling both features silently falls back to non-streaming Layers[i].Forward delegates and reintroduces the unbounded weight working set during training. Either compose checkpoint recomputation over the streaming/materialize path or fail fast until that combination is supported.

Minimum safe guard

         if (segmentSize > 0 && Layers.Count > segmentSize)
         {
+            if (_weightLifetimeConfigured)
+            {
+                throw new NotSupportedException(
+                    "Gradient checkpointing is not yet compatible with weight streaming. " +
+                    "Disable one of them until the training path composes both features.");
+            }
+
             // Cache the layer-forward delegate array so checkpointed training
             // doesn't allocate N closures + a delegate array on every call.
             // Rebuild only when the layer graph changes (structure version

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@src/NeuralNetworks/NeuralNetworkBase.cs` around lines 2669 - 2713, The
checkpointing early-return bypasses the streaming/materialize weight-lifetime
path: before returning the call to
AiDotNet.Tensors.Engines.Autodiff.GradientCheckpointing<T>.Checkpoint, check the
instance flag _weightLifetimeConfigured (from NeuralNetworkBase) and if it's
true fail fast by throwing a clear NotSupportedException indicating
checkpointing cannot be used with configured weight lifetime/streaming;
alternatively implement composition by routing checkpoint recomputation through
the streaming/materialize forward path (i.e. use the same
materialize/streaming-forward logic used elsewhere instead of directly binding
Layers[i].Forward), but until that composition exists prefer the fail-fast
exception so the unsupported combination is explicit.

---

2978-2998: ⚠️ Potential issue | 🟠 Major | ⚡ Quick win

Refresh the weight registry after lazy weights materialize.

ConfigureWeightLifetime() registers only tensors with Length > 0, and this loop never re-registers after Layers[i].Forward(current) allocates real weights for a lazy layer. On explicit streaming for lazy models, later prefetch/materialize passes still miss those tensors, so streaming never fully engages for them.

Suggested fix

         var current = input;
         for (int i = 0; i < Layers.Count; i++)
         {
+            bool registryMayNeedRefresh =
+                Layers[i] is LayerBase<T> layerBefore
+                && layerBefore.GetTrainableParameters().Any(t => t is not null && t.Length == 0);
+
             // Slide the prefetch window forward: by the time we're
             // computing layer i, layers i..i+W-1 should be in the pool;
             // start fetching layer i+W now so it's ready when layer i+1
             // begins. The pre-flight loop above primed i=0; this call
@@
             using (BeginLayerMaterializeScope(i))
             {
                 current = Layers[i].Forward(current);
             }
+
+            if (registryMayNeedRefresh)
+            {
+                RefreshWeightRegistry();
+            }
         }
         return current;
     }

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@src/NeuralNetworks/NeuralNetworkBase.cs` around lines 2978 - 2998,
ConfigureWeightLifetime only registers tensors with Length>0 but lazy layers
allocate real weights during Layers[i].Forward(current), and the current loop
(using BeginLayerMaterializeScope and PrefetchLayerWeights) never re-registers
them, so later streaming passes miss these tensors; fix by invoking the
weight-registration routine after a lazy layer materializes (e.g., call
ConfigureWeightLifetime or a new RegisterLayerWeights method right after
Layers[i].Forward(current) inside the using(BeginLayerMaterializeScope(i))
block), ensuring newly-allocated tensors get added to the weight lifetime
registry so subsequent PrefetchLayerWeights and streaming iterations see them.

src/AiModelBuilder.cs (1)

5467-5470: ⚠️ Potential issue | 🟠 Major

Blocking: Enabled = false is not a guaranteed force-off.

If ctor-time auto-detect already engaged streaming, DisableAutoStreaming() only blocks the lazy retry path; nnBase.IsWeightStreamingActive can still stay true, so an explicit opt-out can build a result that still reports streaming as enabled. This public override should either be enforced or fail fast instead of silently degrading. As per coding guidelines, “Production Readiness (CRITICAL - Flag as BLOCKING): Incomplete features … half-implemented patterns where some code paths work but others silently do nothing.”

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@src/AiModelBuilder.cs` around lines 5467 - 5470, The current branch that
checks _weightStreamingConfig.Enabled only calls nnBase.DisableAutoStreaming()
which may not turn off streaming if ctor-time auto-detect already activated it;
update the public override that checks _weightStreamingConfig.Enabled so that
after calling nnBase.DisableAutoStreaming() you verify
nnBase.IsWeightStreamingActive is false and enforce the opt-out: if a force-off
API exists (e.g., nnBase.ForceDisableWeightStreaming or similar) call it,
otherwise throw a clear InvalidOperationException indicating streaming could not
be disabled; add a small unit test or assertion to cover the path and ensure the
method fails fast instead of silently leaving IsWeightStreamingActive true.

tests/AiDotNet.Tests/IntegrationTests/WeightStreaming/WeightStreamingBuilderIntegrationTests.cs (1)

44-54: ⚠️ Potential issue | 🔴 Critical | ⚡ Quick win

Blocking: ConfigureWeightStreaming_AcceptsAllThreeEnabledStates is assertion-free and can pass on regressions.

At Line 53, this test currently succeeds on “no throw” only. Add unconditional assertions so it fails if fluent behavior/state handling regresses.

Suggested tightening

 [Fact]
 public void ConfigureWeightStreaming_AcceptsAllThreeEnabledStates()
 {
     // Three valid states for the Enabled flag: null (auto-detect),
     // true (force on), false (force off). Pin that all three are
     // accepted by the API.
     var builder = new AiModelBuilder<double, Tensor<double>, Tensor<double>>();
-    builder.ConfigureWeightStreaming(new WeightStreamingConfig { Enabled = null });
-    builder.ConfigureWeightStreaming(new WeightStreamingConfig { Enabled = true });
-    builder.ConfigureWeightStreaming(new WeightStreamingConfig { Enabled = false });
-    // No assertion needed — the test passes if no throw occurs.
+    var r1 = builder.ConfigureWeightStreaming(new WeightStreamingConfig { Enabled = null });
+    var r2 = builder.ConfigureWeightStreaming(new WeightStreamingConfig { Enabled = true });
+    var r3 = builder.ConfigureWeightStreaming(new WeightStreamingConfig { Enabled = false });
+    Assert.Same(builder, r1);
+    Assert.Same(builder, r2);
+    Assert.Same(builder, r3);
 }

As per coding guidelines: “Test Quality (CRITICAL - Flag as BLOCKING) ... Missing assertions.”

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In
`@tests/AiDotNet.Tests/IntegrationTests/WeightStreaming/WeightStreamingBuilderIntegrationTests.cs`
around lines 44 - 54, The test currently only ensures no exception; instead
after each ConfigureWeightStreaming call, add unconditional assertions that the
builder recorded the provided Enabled value (null, true, false) — i.e., call
whatever accessor or produced object exposes the weight-streaming configuration
(for example via a GetWeightStreamingConfig, a Builder property, or by building
the model and inspecting its WeightStreaming/WeightStreamingConfig) and assert
that WeightStreamingConfig.Enabled equals the value passed for each case; keep
the three ConfigureWeightStreaming calls but follow each with a corresponding
Assert.Equal/Assert.IsTrue/Assert.IsNull check against the retrieved config so
regressions in ConfigureWeightStreaming, WeightStreamingConfig, or
AiModelBuilder state handling will fail the test.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@src/AiModelBuilder.cs`:
- Line 3375: BuildStreamingSupervisedAsync currently returns
AiModelResultOptions<T, TInput, TOutput> and does not populate
WeightStreamingReport for builds originating from IStreamingDataLoader, causing
the new telemetry to be omitted; update BuildStreamingSupervisedAsync (and any
streaming-data-loader build paths) to return or produce the options type that
includes WeightStreamingReport and set WeightStreamingReport =
BuildWeightStreamingReport() where the result is constructed so streaming builds
receive the same telemetry as non-streaming paths; ensure all code paths
referenced (BuildStreamingSupervisedAsync, IStreamingDataLoader build paths, and
any factory/producer that returns AiModelResultOptions<T, TInput, TOutput>) are
updated consistently to include the WeightStreamingReport field.
- Around line 5326-5335: The ConfigureWeightStreaming method currently accepts a
WeightStreamingConfig with ThresholdParameters <= 0 which later produces
misleading telemetry; validate the config before assigning
_weightStreamingConfig by checking config is null or config.ThresholdParameters
> 0 and if ThresholdParameters is invalid throw an ArgumentOutOfRangeException
(or ArgumentException) naming ThresholdParameters and the invalid value; update
ApplyWeightStreamingConfig and BuildWeightStreamingReport to assume the field is
valid (or defensively guard), and apply the same validation to the other places
that set _weightStreamingConfig (the other
ConfigureWeightStreaming/weight-streaming setters referenced by
ApplyWeightStreamingConfig/BuildWeightStreamingReport) so invalid external
inputs are rejected up-front.
- Around line 1422-1427: ApplyWeightStreamingConfig() is only called once on the
initial _model, so when the AutoML path reassigns _model (e.g., _model =
bestModel) or the RL AutoML path assigns selectedAgent the weight-streaming
overrides stop applying; ensure ConfigureWeightStreaming /
ApplyWeightStreamingConfig runs after every reassignment of _model by either (a)
invoking ApplyWeightStreamingConfig() immediately after assignments to _model in
AutoML and RL AutoML code paths (mentions: bestModel, selectedAgent) or (b)
centralizing the behavior by adding a single setter/wrapper for the _model field
that assigns the instance and then calls ApplyWeightStreamingConfig() /
ConfigureWeightStreaming(...) so all code paths automatically re-apply overrides
to the actual model used for the rest of the build.

In `@src/NeuralNetworks/NeuralNetworkBase.cs`:
- Around line 2555-2575: The finally block currently marks
_firstForwardCompleted and calls TryAutoEnableWeightStreaming even if
PredictEager(promoted) threw; introduce a local bool (e.g., eagerSuccess =
false), set eagerSuccess = true immediately after PredictEager returns
successfully, and in the finally block only set _firstForwardCompleted = true
and call TryAutoEnableWeightStreaming when eagerSuccess is true (preserving
SetTrainingMode(wasTraining) behavior regardless); update references around
PredictEager, _firstForwardCompleted, and TryAutoEnableWeightStreaming
accordingly.
- Around line 3592-3606: The code currently reads the throwing ParameterCount
property inside a try/catch and returns on exception, which prevents auto-detect
for models > int.MaxValue; replace that block with a non-throwing long
computation (e.g., call or add a helper like ComputeTotalParameterCountLong or
sum the Layers' parameter counts into a long) to obtain paramCount without
invoking the throwing flat API, keep the existing guard for
partially-initialized models (check if layers/InitializeLayers completed before
computing) and do not swallow valid errors; ensure subsequent logic (including
ConfigureWeightLifetime) uses this long paramCount so auto-detect can engage for
very large models.

---

Duplicate comments:
In `@src/AiModelBuilder.cs`:
- Around line 5467-5470: The current branch that checks
_weightStreamingConfig.Enabled only calls nnBase.DisableAutoStreaming() which
may not turn off streaming if ctor-time auto-detect already activated it; update
the public override that checks _weightStreamingConfig.Enabled so that after
calling nnBase.DisableAutoStreaming() you verify nnBase.IsWeightStreamingActive
is false and enforce the opt-out: if a force-off API exists (e.g.,
nnBase.ForceDisableWeightStreaming or similar) call it, otherwise throw a clear
InvalidOperationException indicating streaming could not be disabled; add a
small unit test or assertion to cover the path and ensure the method fails fast
instead of silently leaving IsWeightStreamingActive true.

In `@src/NeuralNetworks/NeuralNetworkBase.cs`:
- Around line 2669-2713: The checkpointing early-return bypasses the
streaming/materialize weight-lifetime path: before returning the call to
AiDotNet.Tensors.Engines.Autodiff.GradientCheckpointing<T>.Checkpoint, check the
instance flag _weightLifetimeConfigured (from NeuralNetworkBase) and if it's
true fail fast by throwing a clear NotSupportedException indicating
checkpointing cannot be used with configured weight lifetime/streaming;
alternatively implement composition by routing checkpoint recomputation through
the streaming/materialize forward path (i.e. use the same
materialize/streaming-forward logic used elsewhere instead of directly binding
Layers[i].Forward), but until that composition exists prefer the fail-fast
exception so the unsupported combination is explicit.
- Around line 2978-2998: ConfigureWeightLifetime only registers tensors with
Length>0 but lazy layers allocate real weights during
Layers[i].Forward(current), and the current loop (using
BeginLayerMaterializeScope and PrefetchLayerWeights) never re-registers them, so
later streaming passes miss these tensors; fix by invoking the
weight-registration routine after a lazy layer materializes (e.g., call
ConfigureWeightLifetime or a new RegisterLayerWeights method right after
Layers[i].Forward(current) inside the using(BeginLayerMaterializeScope(i))
block), ensuring newly-allocated tensors get added to the weight lifetime
registry so subsequent PrefetchLayerWeights and streaming iterations see them.

In
`@tests/AiDotNet.Tests/IntegrationTests/WeightStreaming/WeightStreamingBuilderIntegrationTests.cs`:
- Around line 44-54: The test currently only ensures no exception; instead after
each ConfigureWeightStreaming call, add unconditional assertions that the
builder recorded the provided Enabled value (null, true, false) — i.e., call
whatever accessor or produced object exposes the weight-streaming configuration
(for example via a GetWeightStreamingConfig, a Builder property, or by building
the model and inspecting its WeightStreaming/WeightStreamingConfig) and assert
that WeightStreamingConfig.Enabled equals the value passed for each case; keep
the three ConfigureWeightStreaming calls but follow each with a corresponding
Assert.Equal/Assert.IsTrue/Assert.IsNull check against the retrieved config so
regressions in ConfigureWeightStreaming, WeightStreamingConfig, or
AiModelBuilder state handling will fail the test.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: ASSERTIVE

Plan: Pro

Run ID: 4d51994c-f573-4cf1-b8a7-e37219395211

📥 Commits

Reviewing files that changed from the base of the PR and between 044fc07 and 190801e.

📒 Files selected for processing (5)

• src/AiModelBuilder.cs
• src/Deployment/Configuration/WeightStreamingReport.cs
• src/NeuralNetworks/NeuralNetworkBase.cs
• tests/AiDotNet.Tests/IntegrationTests/WeightStreaming/WeightStreamingBuilderIntegrationTests.cs
• tests/AiDotNet.Tests/IntegrationTests/WeightStreaming/WeightStreamingEndToEndTests.cs

[Copilot]

Pull request overview

Copilot reviewed 11 out of 11 changed files in this pull request and generated 8 comments.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

Pull request overview

Copilot reviewed 167 out of 167 changed files in this pull request and generated 6 comments.

[coderabbitai]

Actionable comments posted: 13

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (11)

src/NeuralNetworks/Layers/DigitCapsuleLayer.cs (3)

890-901: ⚠️ Potential issue | 🟠 Major | ⚡ Quick win

SetParameters breaks streaming registration by creating a new tensor.

Same issue as UpdateParameters—Line 898 creates a fresh Tensor<T> and assigns it to _weights, discarding the lazy allocation registration. This would break streaming when deserializing or restoring model weights.

Write directly into the existing tensor:

Suggested fix

 public override void SetParameters(Vector<T> parameters)
 {
     if (parameters.Length != _weights.Length)
     {
         throw new ArgumentException($"Expected {_weights.Length} parameters, but got {parameters.Length}");
     }

-    // Write parameters directly into a new mutable tensor
-    _weights = new Tensor<T>(_weights._shape);
-    for (int i = 0; i < parameters.Length; i++)
-        _weights[i] = parameters[i];
+    // Write parameters directly into the existing lazy-allocated tensor
+    var span = _weights.AsWritableSpan();
+    for (int i = 0; i < parameters.Length; i++)
+        span[i] = parameters[i];
 }

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@src/NeuralNetworks/Layers/DigitCapsuleLayer.cs` around lines 890 - 901,
SetParameters currently replaces the existing _weights by new
Tensor<T>(_weights._shape), which breaks lazy allocation/streaming registration
(same bug as UpdateParameters); instead, do not reassign _weights—validate
parameters.Length matches _weights.Length and write each value directly into the
existing _weights tensor (e.g., loop assigning to _weights[i]) so the original
Tensor<T> instance and its streaming/lazy-registration remain intact.

---

804-810: ⚠️ Potential issue | 🟠 Major | ⚡ Quick win

UpdateParameters breaks streaming registration by reassigning _weights.

The method computes a new tensor via _weights.Subtract(...) and assigns it back to _weights, discarding the lazy allocation registration established in EnsureInitialized. Every training iteration would break streaming.

Apply the same in-place update pattern used in InitializeParameters:

Suggested fix

 public override void UpdateParameters(T learningRate)
 {
     if (_weightsGradient == null)
         throw new InvalidOperationException("Backward pass must be called before updating parameters.");

-    _weights = _weights.Subtract(_weightsGradient.Multiply(learningRate));
+    var delta = _weightsGradient.Multiply(learningRate);
+    var updated = _weights.Subtract(delta);
+    updated.AsSpan().CopyTo(_weights.AsWritableSpan());
 }

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@src/NeuralNetworks/Layers/DigitCapsuleLayer.cs` around lines 804 - 810,
UpdateParameters currently replaces the tensor referenced by _weights with a new
result from _weights.Subtract(...), breaking the streaming/lazy-allocation
registration created by EnsureInitialized; instead perform the update in-place
like InitializeParameters does by applying the gradient-scaled subtraction into
the existing _weights buffer (use the same in-place method used in
InitializeParameters), ensuring you read _weightsGradient, multiply by
learningRate, and subtract it into the existing _weights storage rather than
reassigning the field so streaming remains registered.

---

292-304: ⚠️ Potential issue | 🔴 Critical

BLOCKING: Eager constructor must use AllocateLazyWeight to enable streaming support for trainable weights

The _weights field is marked [TrainableParameter], yet the eager constructor (line 300) uses new Tensor<T>(...) instead of AllocateLazyWeight. The lazy path correctly uses AllocateLazyWeight at line 428.

This creates an asymmetry: the eager constructor bypasses the streaming pool entirely, while the lazy path registers properly. Since _weights is explicitly trainable and can grow to multi-GB scale (especially in multi-class scenarios), both paths must register through the streaming pool.

Fix: Change line 300 to:

_weights = AllocateLazyWeight([inputCapsules, numClasses, inputCapsuleDimension, outputCapsuleDimension]);

Then ensure InitializeParameters() copies values into the existing tensor in place (consistent with the pattern documented in PrimaryCapsuleLayer lines 477–478 and RecurrentLayer lines 961–962).

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@src/NeuralNetworks/Layers/DigitCapsuleLayer.cs` around lines 292 - 304, The
eager constructor in DigitCapsuleLayer currently assigns _weights with new
Tensor<T>(...) which bypasses the streaming pool; replace that assignment with
_weights = AllocateLazyWeight([inputCapsules, numClasses, inputCapsuleDimension,
outputCapsuleDimension]) so the trainable _weights is registered for streaming,
and then update InitializeParameters() to fill values into the already-allocated
_weights in-place (same pattern used by PrimaryCapsuleLayer and RecurrentLayer)
rather than replacing the tensor instance.

src/NeuralNetworks/Layers/LocallyConnectedLayer.cs (1)

1017-1019: ⚠️ Potential issue | 🟠 Major | ⚡ Quick win

Tensor identity broken by GPU weight sync.

Same issue as UpdateParameters — these lines replace _weights and _biases field references with the GPU tensor objects, orphaning the original registered tensors:

_weights = _gpuWeights;
_biases = _gpuBiases;

If weight streaming is engaged, the registry will hold stale references after any GPU-based parameter update.

🐛 Proposed fix

     // Sync back to CPU tensors for compatibility
-    _weights = _gpuWeights;
-    _biases = _gpuBiases;
+    // Copy GPU data back to CPU tensors in-place to preserve streaming registry
+    _gpuWeights.Data.Span.CopyTo(_weights.Data.Span);
+    _gpuBiases.Data.Span.CopyTo(_biases.Data.Span);
+    
+    // Notify engine that parameters have changed
+    Engine.InvalidatePersistentTensor(_weights);
+    Engine.InvalidatePersistentTensor(_biases);

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@src/NeuralNetworks/Layers/LocallyConnectedLayer.cs` around lines 1017 - 1019,
The current lines in LocallyConnectedLayer.cs (_weights = _gpuWeights; _biases =
_gpuBiases;) replace the registered CPU tensor references with GPU tensor
objects and orphan the original tensors; instead, copy the GPU contents back
into the existing registered CPU tensors (e.g., call the tensors'
CopyFrom/CopyTo or equivalent methods on the existing _weights and _biases) and
only assign new tensors if the registered ones are null, ensuring you do not
reassign the _weights/_biases fields to _gpuWeights/_gpuBiases; mirror the same
pattern used to fix UpdateParameters so the registry keeps valid references.

src/NeuralNetworks/Layers/Conv3DLayer.cs (1)

757-789: ⚠️ Potential issue | 🟠 Major | ⚡ Quick win

Keep the streaming-aware tensors when loading parameters.

Lines 783-785 replace _kernels and _biases with fresh Tensor<T> instances after the deferred-shape path has resolved the layer. That throws away the lazy/registered tensors created for streaming, so a model loaded through SetParameters() can silently fall back to non-streamed weights for this layer. Populate the existing tensors in place instead of reassigning the fields.

Proposed fix

         int index = 0;
-        _kernels = new Tensor<T>(_kernels._shape, parameters.Slice(index, _kernels.Length));
+        var kernelSpan = _kernels.AsWritableSpan();
+        for (int i = 0; i < _kernels.Length; i++)
+        {
+            kernelSpan[i] = parameters[index + i];
+        }
         index += _kernels.Length;
-        _biases = new Tensor<T>(_biases._shape, parameters.Slice(index, _biases.Length));
+        var biasSpan = _biases.AsWritableSpan();
+        for (int i = 0; i < _biases.Length; i++)
+        {
+            biasSpan[i] = parameters[index + i];
+        }

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@src/NeuralNetworks/Layers/Conv3DLayer.cs` around lines 757 - 789,
SetParameters replaces streaming-aware tensors (_kernels, _biases) with new
Tensor<T> instances which discards lazy/registered streaming state; instead,
detect when _kernels/_biases already exist and copy the parameter slices into
those existing tensors in-place (e.g., use an existing tensor
copy/assign/SetData method or a buffer copy from parameters.Slice(...)) and only
allocate new Tensor<T> when the fields are null; preserve the existing tensor
instances so streaming registration is kept, then call
Engine.InvalidatePersistentTensor on the existing tensors after updating their
contents and keep the same validation and length checks in SetParameters.

src/LoRA/Adapters/LoRETTAAdapter.cs (1)

127-148: ⚠️ Potential issue | 🟡 Minor | ⚡ Quick win

Carry the 64-bit fix through the reporting path.

ParameterCount is now safe, but GetParameterEfficiencyMetrics() still narrows fullParams and ttParams back to int. Large adapters will report truncated or negative numbers even though the main parameter path was widened here.

Suggested fix

-        int fullParams = inputSize * outputSize;
-        int ttParams = (int)(ParameterCount - (_freezeBaseLayer ? 0 : _baseLayer.ParameterCount));
-        int equivalentLoRAParams = (inputSize + outputSize) * TTRank;
+        long fullParams = (long)inputSize * outputSize;
+        long ttParams = ParameterCount - (_freezeBaseLayer ? 0L : _baseLayer.ParameterCount);
+        long equivalentLoRAParams = (long)(inputSize + outputSize) * TTRank;

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@src/LoRA/Adapters/LoRETTAAdapter.cs` around lines 127 - 148,
GetParameterEfficiencyMetrics() still uses int for fullParams and ttParams
causing truncation; update that method to use long (Int64) for fullParams,
ttParams and any intermediate multiplications (use long casts where multiplying
_ttRanks and _coreShapes) and ensure you sum _baseLayer.ParameterCount (already
long) into a long variable; also update any formatting/returns that assume int
to handle long values so reported metrics match ParameterCount and avoid
negative/wrapped numbers.

src/NeuralNetworks/Layers/SpiralConvLayer.cs (1)

497-554: ⚠️ Potential issue | 🟠 Major | 🏗️ Heavy lift

Lazy weight identity is still lost after initialization.

This path now preserves the lazy tensors during init, but UpdateParameters(), SetParameters(), and Deserialize() below still replace _weights and _biases with new Tensor<T> instances. After the first optimizer step or parameter load, the registry will still be tracking the old tensors, so streaming/prefetching stops applying to the live weights.

Suggested fix

-        _weights = Engine.TensorSubtract(_weights, scaledWeightGrad);
+        var updatedWeights = Engine.TensorSubtract(_weights, scaledWeightGrad);
+        updatedWeights.AsSpan().CopyTo(_weights.AsWritableSpan());

-        _biases = Engine.TensorSubtract(_biases, scaledBiasGrad);
+        var updatedBiases = Engine.TensorSubtract(_biases, scaledBiasGrad);
+        updatedBiases.AsSpan().CopyTo(_biases.AsWritableSpan());

Apply the same in-place-copy pattern in SetParameters() and Deserialize() so the lazy-allocated tensors remain the authoritative objects for the full layer lifecycle.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@src/NeuralNetworks/Layers/SpiralConvLayer.cs` around lines 497 - 554, The
bug: UpdateParameters(), SetParameters(), and Deserialize() replace the fields
_weights and _biases with new Tensor<T> instances which breaks the
previously-registered lazy tensors; fix by changing those methods to copy
incoming tensor data into the existing lazy-allocated tensors in-place (use
incomingTensor.AsSpan().CopyTo(_weights.AsWritableSpan()) /
.CopyTo(_biases.AsWritableSpan())), only reallocate if shapes differ, and do not
re-register new Tensor objects so the original registered lazy tensors remain
authoritative for streaming/prefetching.

src/NeuralNetworks/Layers/DenseLayer.cs (1)

1354-1389: ⚠️ Potential issue | 🔴 Critical | ⚡ Quick win

Blocking: reject incompatible checkpoint lengths before copying into initialized tensors.

This path currently accepts unresolved, truncated, or shape-mismatched payloads and then copies only Math.Min(...) elements, leaving the remainder at fresh init values. That silently turns a bad checkpoint into a hybrid model instead of failing deterministically at deserialize time. Validate wLen/bLen before EnsureInitialized() and require exact matches once shapes are resolved.

Proposed fix

 int wLen = reader.ReadInt32();
 if (!IsShapeResolved)
 {
     int outputSize = OutputShape[0];
-    if (outputSize > 0 && wLen > 0 && wLen % outputSize == 0)
-    {
-        int inferredInput = wLen / outputSize;
-        ResolveFromShape(new[] { inferredInput });
-    }
+    if (outputSize <= 0 || wLen <= 0 || wLen % outputSize != 0)
+    {
+        throw new InvalidDataException(
+            $"Cannot resolve DenseLayer shape from serialized weight length {wLen} and output size {outputSize}.");
+    }
+
+    ResolveFromShape(new[] { wLen / outputSize });
 }
 EnsureInitialized();
+if (wLen != _weights.Length)
+{
+    throw new InvalidDataException(
+        $"Serialized DenseLayer weight length {wLen} does not match expected {_weights.Length}.");
+}
 // Read weights IN PLACE to preserve engine's persistent tensor reference
 var wSpan = _weights.Data.Span;
-for (int i = 0; i < Math.Min(wLen, _weights.Length); i++)
+for (int i = 0; i < wLen; i++)
     wSpan[i] = NumOps.FromDouble(reader.ReadDouble());
-// Skip any extra values if serialized layer was bigger
-for (int i = _weights.Length; i < wLen; i++)
-    reader.ReadDouble();

 // Read biases IN PLACE
 int bLen = reader.ReadInt32();
+if (bLen != _biases.Length)
+{
+    throw new InvalidDataException(
+        $"Serialized DenseLayer bias length {bLen} does not match expected {_biases.Length}.");
+}
 var bSpan = _biases.Data.Span;
-for (int i = 0; i < Math.Min(bLen, _biases.Length); i++)
+for (int i = 0; i < bLen; i++)
     bSpan[i] = NumOps.FromDouble(reader.ReadDouble());
-for (int i = _biases.Length; i < bLen; i++)
-    reader.ReadDouble();

As per coding guidelines, missing error handling at system boundaries is blocking.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@src/NeuralNetworks/Layers/DenseLayer.cs` around lines 1354 - 1389, The
deserializer currently allows mismatched checkpoint sizes and silently partially
copies values into initialized tensors; update the Deserialize path so that
before calling EnsureInitialized() you validate wLen and bLen against resolved
shapes (use OutputShape and IsShapeResolved and, if unresolved, infer input as
you already do with ResolveFromShape), and after ResolveFromShape require exact
matches between wLen and _weights.Length and between bLen and _biases.Length; if
sizes differ throw a descriptive exception (e.g. InvalidDataException) instead
of using Math.Min and copying partial data; keep the existing
Engine.InvalidatePersistentTensor(_weights) and
Engine.InvalidatePersistentTensor(_biases) behavior after successful validation
and copy.

src/LoRA/Adapters/ChainLoRAAdapter.cs (1)

331-366: ⚠️ Potential issue | 🔴 Critical | 🏗️ Heavy lift

Frozen adapters still affect inference, so excluding them from packed state makes round-trips lossy.

FreezeActiveAdapter() does not merge anything into the base layer, and Forward() still sums every adapter. But this count/allocation path now drops frozen adapters from Parameters and ParameterGradients. After a freeze, clone/save/load can lose still-active adapter weights and change predictions. Persisted state needs to include frozen adapters even if trainable-state bookkeeping does not.

As per coding guidelines, incomplete features where some code paths silently drop active state are blocking.

Also applies to: 543-579

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@src/LoRA/Adapters/ChainLoRAAdapter.cs` around lines 331 - 366, ParameterCount
currently omits adapters that are frozen (or marked merged) causing saved/loaded
packed state to drop weights used at inference; update the calculation in the
ParameterCount getter (and the analogous ParameterGradients code paths around
the other noted block) to always include adapter parameter counts for adapters
present in _adapterChain regardless of _mergedStatus or freeze state, while
still excluding only truly absent/removed adapters; reference _adapterChain,
_mergedStatus, _freezeBaseLayer, _baseLayer, _currentParameterCount,
_chainLength, FreezeActiveAdapter(), and Forward() to locate the related logic
and ensure the same inclusion rule is applied where state is
serialized/persisted so frozen-but-active adapters are retained.

src/NeuralNetworks/Layers/RecurrentLayer.cs (1)

832-885: ⚠️ Potential issue | 🟠 Major | ⚡ Quick win

SetParameters replaces the lazy tensors instead of populating them.

Once the layer has been initialized, this method allocates new tensors for all three parameter blocks and re-registers them. That throws away the lazy/streamable instances created in EnsureInitialized, so loaded models stop using the registered weights the streaming path warmed.

Minimal fix

-        _inputWeights = new Tensor<T>(_inputWeights._shape);
         for (int i = 0; i < inputWeightsSize; i++)
             _inputWeights[i] = parameters[idx++];

-        _hiddenWeights = new Tensor<T>(_hiddenWeights._shape);
         for (int i = 0; i < hiddenWeightsSize; i++)
             _hiddenWeights[i] = parameters[idx++];

-        _biases = new Tensor<T>(_biases._shape);
         for (int i = 0; i < _biases.Length; i++)
             _biases[i] = parameters[idx++];

-        RegisterTrainableParameter(_inputWeights, PersistentTensorRole.Weights);
-        RegisterTrainableParameter(_hiddenWeights, PersistentTensorRole.Weights);
-        RegisterTrainableParameter(_biases, PersistentTensorRole.Biases);

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@src/NeuralNetworks/Layers/RecurrentLayer.cs` around lines 832 - 885,
SetParameters currently discards the lazy/streamable tensors by allocating new
Tensor<T> instances for _inputWeights, _hiddenWeights and _biases and
re-registering them; instead, populate the existing tensors created by
EnsureInitialized so the streaming path keeps using the same registered
instances. Modify RecurrentLayer.SetParameters to validate parameters.Length
against the expected total, then copy values into the existing _inputWeights,
_hiddenWeights and _biases buffers (e.g., assign into their indices) rather than
newing them up; only call RegisterTrainableParameter when a tensor was actually
newly created (or remove the duplicate registrations), and preserve the original
tensor shapes (_inputWeights._shape, _hiddenWeights._shape, _biases._shape) when
copying.

src/NeuralNetworks/Layers/GRULayer.cs (1)

1476-1520: ⚠️ Potential issue | 🔴 Critical | ⚡ Quick win

Blocking: reject any parameter vector that does not exactly match the GRU layout.

When the lazy inference at Lines 1481-1509 cannot derive a valid shape, this still falls through into the bulk copies at Lines 1512-1520. On a fresh lazy layer those targets are zero-length placeholders, so a bad checkpoint can be accepted as a silent no-op; on an already-resolved layer, extra trailing values are silently ignored. SetParameters needs an exact-length guard before copying.

Proposed fix

 public override void SetParameters(Vector<T> parameters)
 {
     // Lazy ctor: if shape isn't resolved (placeholder W/U/b tensors
     // with Length 0), infer inputSize from the param vector. Layout:
     //   3*W[hiddenSize, inputSize] + 3*U[hiddenSize, hiddenSize] + 3*b[hiddenSize]
     //   = 3*hiddenSize*(inputSize + hiddenSize + 1)
     // → inputSize = total/(3*hiddenSize) - hiddenSize - 1.
     if (!IsShapeResolved && _hiddenSize > 0)
     {
         int divisor = 3 * _hiddenSize;
-        if (parameters.Length % divisor == 0)
-        {
-            int candidateInput = parameters.Length / divisor - _hiddenSize - 1;
-            if (candidateInput > 0)
-            {
-                // The divisibility check alone can pass for malformed
-                // vectors that happen to land on a multiple of
-                // 3*hiddenSize but don't actually correspond to any
-                // valid (inputSize, hiddenSize) GRU layout. Reverify
-                // the round-trip count exactly so we reject those
-                // before they corrupt the resolved shape.
-                long expectedCount = (long)_hiddenSize * candidateInput * 3 +
-                                     (long)_hiddenSize * _hiddenSize * 3 +
-                                     (long)_hiddenSize * 3;
-                if (expectedCount != parameters.Length)
-                {
-                    throw new ArgumentException(
-                        $"Parameter vector length {parameters.Length} does not match GRU layout " +
-                        $"for inferred inputSize={candidateInput}, hiddenSize={_hiddenSize} " +
-                        $"(expected {expectedCount}).",
-                        nameof(parameters));
-                }
-                ResolveFromShape(new[] { candidateInput });
-            }
-        }
+        if (parameters.Length % divisor != 0)
+        {
+            throw new ArgumentException(
+                $"Parameter vector length {parameters.Length} does not match a GRU layout for hiddenSize={_hiddenSize}.",
+                nameof(parameters));
+        }
+
+        int candidateInput = parameters.Length / divisor - _hiddenSize - 1;
+        long expectedCount = (long)_hiddenSize * candidateInput * 3 +
+                             (long)_hiddenSize * _hiddenSize * 3 +
+                             (long)_hiddenSize * 3;
+        if (candidateInput <= 0 || expectedCount != parameters.Length)
+        {
+            throw new ArgumentException(
+                $"Parameter vector length {parameters.Length} does not match a GRU layout for hiddenSize={_hiddenSize}.",
+                nameof(parameters));
+        }
+
+        ResolveFromShape([candidateInput]);
     }
+
+    EnsureInitialized();
+    int expectedLength = ParameterCountHelper.ToFlatVectorSize(ParameterCount);
+    if (parameters.Length != expectedLength)
+    {
+        throw new ArgumentException(
+            $"Expected {expectedLength} parameters for GRULayer(hiddenSize={_hiddenSize}, inputSize={_inputSize}), got {parameters.Length}.",
+            nameof(parameters));
+    }
+
     // Bulk copy from parameter vector into tensor storage — avoids per-element SetFlat calls
     int idx = 0;
     parameters.Slice(idx, _Wz.Length).AsSpan().CopyTo(_Wz.Data.Span); idx += _Wz.Length;
     parameters.Slice(idx, _Wr.Length).AsSpan().CopyTo(_Wr.Data.Span); idx += _Wr.Length;
     parameters.Slice(idx, _Wh.Length).AsSpan().CopyTo(_Wh.Data.Span); idx += _Wh.Length;

As per coding guidelines, “missing validation of external inputs” and “incomplete features” are blocking issues.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@src/NeuralNetworks/Layers/GRULayer.cs` around lines 1476 - 1520, The code
allows malformed parameter vectors to slip through to the bulk CopyTo calls; fix
SetParameters by validating the parameter vector length exactly matches the GRU
layout before any copying. After the lazy-inference block (the code using
IsShapeResolved, _hiddenSize and ResolveFromShape) compute the expectedCount
from the now-resolved input size and _hiddenSize (expectedCount = 3 *
_hiddenSize * (inputSize + _hiddenSize + 1)), verify parameters.Length ==
expectedCount, and if not throw an ArgumentException naming parameters; do this
prior to the sequence of parameters.Slice(...).CopyTo(...) operations to prevent
silent no-ops or truncation.

♻️ Duplicate comments (8)

src/NeuralNetworks/Layers/DecoderLayer.cs (1)

152-160: ⚠️ Potential issue | 🔴 Critical | ⚡ Quick win

Validate that SetParameters consumes the entire vector.

This still accepts trailing parameter data silently. A stale or incompatible checkpoint can therefore “deserialize” without error while leaving extra values unread. Please assert idx == parameters.Length after the last Set(...). As per coding guidelines, production-ready code must not have “missing validation of external inputs”.

Suggested fix

 Set(_selfAttention); Set(_crossAttention); Set(_feedForward1); Set(_feedForward2);
 Set(_norm1); Set(_norm2); Set(_norm3);
+if (idx != parameters.Length)
+{
+    throw new InvalidOperationException(
+        $"DecoderLayer expected to consume {idx} parameters but received {parameters.Length}.");
+}

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@src/NeuralNetworks/Layers/DecoderLayer.cs` around lines 152 - 160, After
calling Set(_selfAttention); Set(_crossAttention); Set(_feedForward1);
Set(_feedForward2); Set(_norm1); Set(_norm2); Set(_norm3); validate that the
entire parameters span was consumed by checking that idx == parameters.Length
and fail fast if not; add an explicit check after the last Set(...) that throws
a clear InvalidOperationException (or Debug.Assert in debug builds) with a
message referencing the expected and actual counts (use idx and
parameters.Length) so stale/incompatible checkpoints cannot be silently
accepted.

src/LoRA/Adapters/NOLAAdapter.cs (1)

225-228: ⚠️ Potential issue | 🟠 Major | ⚡ Quick win

Large unfrozen adapters still narrow the backing parameter buffer.

ParameterCount is 64-bit now, but the constructor still sizes Parameters with (int)(_baseLayer.ParameterCount + nolaParams), and UpdateCoefficientsFromParameters() still narrows _baseLayer.ParameterCount before unpacking. Large unfrozen adapters can still overflow or fail before ParameterCountHelper ever validates the size.

Suggested fix

-        int nolaParams = 2 * _numBasis;
-        Parameters = new Vector<T>(_freezeBaseLayer ? nolaParams : (int)(_baseLayer.ParameterCount + nolaParams));
+        int nolaParams = checked(2 * _numBasis);
+        Parameters = new Vector<T>(ParameterCountHelper.ToFlatVectorSize(ParameterCount));
         UpdateParametersFromCoefficients();

-            int baseParamCount = checked((int)_baseLayer.ParameterCount);
+            int baseParamCount = ParameterCountHelper.ToFlatVectorSize(_baseLayer.ParameterCount);
             Vector<T> baseParams = new Vector<T>(baseParamCount);

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@src/LoRA/Adapters/NOLAAdapter.cs` around lines 225 - 228, The constructor and
UpdateCoefficientsFromParameters still narrow 64-bit counts into ints causing
overflow for large unfrozen adapters; update the NOLAAdapter constructor to
validate total size with ParameterCountHelper using long arithmetic (baseCount +
nolaParams) before allocation, avoid casting the computed total to int when
sizing the Parameters buffer (use a long-aware allocation or split into chunks
if the runtime needs ints), and change UpdateCoefficientsFromParameters to use
long indexes/lengths when reading/unpacking from _baseLayer.ParameterCount and
_numBasis so you never cast/truncate the long parameter counts to int. Ensure
all references in this file to casting expressions like
(int)(_baseLayer.ParameterCount + nolaParams) are removed and replaced with
long-safe checks and buffer handling, and call ParameterCountHelper early to
fail fast on oversized requests.

src/NeuralNetworks/Layers/GatedLinearUnitLayer.cs (1)

737-778: ⚠️ Potential issue | 🟠 Major

SetParameters still drops the lazy/streaming tensors.

After the new shape-inference block, this method recreates all four parameter tensors with new Tensor<T>(...). That discards the handles allocated and registered in OnFirstForward, so a loaded model falls back to eager resident tensors instead of the streaming-aware ones.

Minimal fix

-        _linearWeights = new Tensor<T>(_linearWeights._shape, parameters.Slice(index, linearWeightsSize));
-        index += linearWeightsSize;
-        _gateWeights = new Tensor<T>(_gateWeights._shape, parameters.Slice(index, gateWeightsSize));
-        index += gateWeightsSize;
-        _linearBias = new Tensor<T>(_linearBias._shape, parameters.Slice(index, _linearBias.Length));
-        index += _linearBias.Length;
-        _gateBias = new Tensor<T>(_gateBias._shape, parameters.Slice(index, _gateBias.Length));
+        for (int i = 0; i < linearWeightsSize; i++) _linearWeights[i] = parameters[index++];
+        for (int i = 0; i < gateWeightsSize; i++) _gateWeights[i] = parameters[index++];
+        for (int i = 0; i < _linearBias.Length; i++) _linearBias[i] = parameters[index++];
+        for (int i = 0; i < _gateBias.Length; i++) _gateBias[i] = parameters[index++];

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@src/NeuralNetworks/Layers/GatedLinearUnitLayer.cs` around lines 737 - 778,
SetParameters currently replaces the existing parameter tensor objects
(_linearWeights, _gateWeights, _linearBias, _gateBias) with new Tensor<T>
instances, discarding the streaming/resident handles created in OnFirstForward;
instead, update the existing tensors' storage with slices from the incoming
parameters (or only allocate new Tensor objects when the existing fields are
null/uninitialized) and then call Engine.InvalidatePersistentTensor on those
existing tensors; change the assignments that use new Tensor<T>(...) to write
into the preallocated tensors (e.g., call a SetData/ReplaceBuffer-style method
or copy into their buffer using parameters.Slice ranges) so the streaming-aware
handles are preserved.

src/NeuralNetworks/Layers/SparseLinearLayer.cs (1)

484-535: ⚠️ Potential issue | 🟠 Major

Don't replace _weights with a new SparseTensor<T> after registration.

Both paths assign a fresh sparse tensor to _weights. The trainable-parameter registry still tracks the constructor-time instance, so post-load/training can read one tensor and update another. Please keep the registered sparse object alive and either mutate its backing storage in place or add a registry-aware sparse swap helper instead of direct reassignment.

Also applies to: 550-569

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@src/NeuralNetworks/Layers/SparseLinearLayer.cs` around lines 484 - 535, The
Deserialize method currently replaces the registered _weights field with a new
SparseTensor<T>, which breaks the trainable-parameter registry; instead, avoid
reassigning _weights and either (a) mutate the existing registered
SparseTensor<T> instance’s internal storage in place (overwrite its RowIndices,
ColumnIndices and Values buffers) when the sparsity pattern differs, or (b) call
a registry-aware swap helper to atomically replace the backing arrays while
keeping the original SparseTensor<T> object identity; update the branches in
Deserialize (and the similar block around lines 550-569) to fill or resize the
existing _weights.RowIndices, _weights.ColumnIndices and _weights.Values
buffers, validate their lengths against nnz, and only then call
base.Deserialize(reader)/SetParameters so the registry continues to reference
the same SparseTensor<T> instance.

src/Models/Results/AiModelResult.cs (1)

2985-2994: ⚠️ Potential issue | 🟠 Major | ⚡ Quick win

DeepCopy() still drops the newly preserved telemetry/config.

WithParameters() now carries InterpretabilityOptions and WeightStreamingReport, but DeepCopy() still omits both. That means a cloned result can silently lose explainability settings and the new weight-streaming telemetry even though the other reconstruction paths now preserve them.

💡 Minimal follow-up in DeepCopy()

         var options = new AiModelResultOptions<T, TInput, TOutput>
         {
             OptimizationResult = clonedOptimizationResult,
             PreprocessingInfo = PreprocessingInfo,
             BiasDetector = BiasDetector,
             FairnessEvaluator = FairnessEvaluator,
+            InterpretabilityOptions = InterpretabilityOptions,
             RagRetriever = RagRetriever,
             RagReranker = RagReranker,
             RagGenerator = RagGenerator,
             QueryProcessors = QueryProcessors,
             LoRAConfiguration = LoRAConfiguration,
@@
             HyperparameterTrialId = HyperparameterTrialId,
             Hyperparameters = Hyperparameters,
-            TrainingMetricsHistory = TrainingMetricsHistory
+            TrainingMetricsHistory = TrainingMetricsHistory,
+            WeightStreamingReport = WeightStreamingReport
         };

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@src/Models/Results/AiModelResult.cs` around lines 2985 - 2994, DeepCopy()
currently omits the newly preserved fields InterpretabilityOptions and
WeightStreamingReport introduced for WithParameters(), causing clones to lose
explainability settings and streaming telemetry; update the DeepCopy()
implementation in AiModelResult to copy these fields into the returned clone
(ensure you perform a deep copy/clone of InterpretabilityOptions if it is a
mutable object and copy or clone WeightStreamingReport as appropriate), matching
the same preservation semantics used by WithParameters() so cloned results
retain both InterpretabilityOptions and WeightStreamingReport.

src/AiModelBuilder.cs (1)

5531-5535: ⚠️ Potential issue | 🟠 Major | ⚡ Quick win

Blocking: Enabled = false is still only best-effort.

If ctor-time auto-detect already engaged streaming, this branch just disables future auto-detect and returns with streaming still active. That leaves a public “force off” override silently unenforceable.

Suggested fix

         if (_weightStreamingConfig.Enabled == false)
         {
             nnBase.DisableAutoStreaming();
+            if (nnBase.IsWeightStreamingActive)
+            {
+                throw new NotSupportedException(
+                    "ConfigureWeightStreaming(Enabled=false) cannot disable weight streaming once it has already been engaged for this model instance.");
+            }
             return;
         }

As per coding guidelines: “Production Readiness (CRITICAL - Flag as BLOCKING): Incomplete features … half-implemented patterns where some code paths work but others silently do nothing.”

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@src/AiModelBuilder.cs` around lines 5531 - 5535, The current branch only
calls nnBase.DisableAutoStreaming() when _weightStreamingConfig.Enabled is
false, leaving streaming active if ctor auto-detect already started it; change
the branch so it not only disables future auto-detect via
nnBase.DisableAutoStreaming() but also forcibly stops any currently active
streaming session (e.g., call nnBase.StopStreaming() or
nnBase.ForceDisableStreaming()) so the "Enabled = false" override is always
enforced; if nnBase lacks such a method, add a clear API
(StopStreaming/ForceDisableStreaming) and invoke it here and update related unit
tests to cover ctor-auto-detect + config=false scenarios.

src/NeuralNetworks/NeuralNetworkBase.cs (1)

2697-2741: ⚠️ Potential issue | 🟠 Major | ⚡ Quick win

Checkpointing path still bypasses streaming in training

Line 2697 returns from the checkpoint branch before Line 2757’s _weightLifetimeConfigured branch, so enabling both features silently disables streaming during ForwardForTraining. That reintroduces the unbounded weight working set for checkpointed training.

Suggested minimal safe guard

 if (segmentSize > 0 && Layers.Count > segmentSize)
 {
+    if (_weightLifetimeConfigured)
+    {
+        throw new NotSupportedException(
+            "Gradient checkpointing is not yet compatible with weight streaming in ForwardForTraining. " +
+            "Disable one of them until composed execution is implemented.");
+    }
+
     if (_checkpointLayerFunctions is null
         || _checkpointFunctionsVersion != _layerStructureVersion)
     {

Also applies to: 2757-2769

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@src/NeuralNetworks/NeuralNetworkBase.cs` around lines 2697 - 2741, The
checkpoint branch currently returns immediately after calling
GradientCheckpointing<T>.Checkpoint (using _checkpointLayerFunctions), which
skips the later _weightLifetimeConfigured streaming logic and thus disables
streaming for checkpointed training; fix by ensuring the
weight-lifetime/streaming branch (_weightLifetimeConfigured) runs for
checkpointed forwards as well — either move the _weightLifetimeConfigured
handling to before the checkpoint return, or have the checkpoint result flow
back through the same streaming/postprocessing path (e.g., call
GradientCheckpointing<T>.Checkpoint into a local result and then apply the
existing _weightLifetimeConfigured streaming/cleanup logic before returning) so
ForwardForTraining honors streaming when both features are enabled.

tests/AiDotNet.Tests/IntegrationTests/WeightStreaming/WeightStreamingEndToEndTests.cs (1)

249-256: ⚠️ Potential issue | 🔴 Critical | ⚡ Quick win

Assert actual eviction, not just “some pool activity.”

resident > 0L || report.EvictionCount > 0 can still pass with EvictionCount == 0, which does not prove the “forces eviction” behavior this test claims.

Suggested fix

 var report = AiDotNet.Tensors.LinearAlgebra.WeightRegistry.GetStreamingReport();
-Assert.True(resident > 0L || report.EvictionCount > 0,
-    $"TightPoolBudget test passed the budget bound but produced no "
-    + $"positive streaming activity (ResidentBytes={resident}, "
-    + $"EvictionCount={report.EvictionCount}). This means registration "
-    + "never reached the streaming pool and the test was effectively a "
-    + "no-op. Verify tensor.Lifetime = Streaming and that the pool "
-    + "config was applied before RegisterWeight ran.");
+Assert.True(report.EvictionCount > 0,
+    $"Expected eviction under 1024-byte budget, but EvictionCount={report.EvictionCount} "
+    + $"(ResidentBytes={resident}). The test did not prove eviction occurred.");

As per coding guidelines: “Test Quality (CRITICAL - Flag as BLOCKING) … Missing assertions” and “Assert specific expected values, not just non-null.”

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In
`@tests/AiDotNet.Tests/IntegrationTests/WeightStreaming/WeightStreamingEndToEndTests.cs`
around lines 249 - 256, The test currently allows passing when resident > 0 even
if no eviction occurred; change the assertion to require an actual eviction by
asserting report.EvictionCount > 0 (keep the resident value in the failure
message for diagnostics). Locate the call to
AiDotNet.Tensors.LinearAlgebra.WeightRegistry.GetStreamingReport() (symbol:
GetStreamingReport), the local variable report and the resident variable in the
TightPoolBudget test, and replace the lax assertion (resident > 0L ||
report.EvictionCount > 0) with a strict assertion that report.EvictionCount > 0,
preserving the formatted failure string showing ResidentBytes and EvictionCount.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: ASSERTIVE

Plan: Pro

Run ID: a1fb8e2e-c0e7-4b2c-8c2a-4b4d9a0de98f

📥 Commits

Reviewing files that changed from the base of the PR and between 110463a and 5fa4c44.

📒 Files selected for processing (46)

• src/AiModelBuilder.cs
• src/Deployment/Configuration/WeightStreamingReport.cs
• src/Helpers/ParameterCountHelper.cs
• src/Interfaces/IAiModelBuilder.cs
• src/LoRA/Adapters/ChainLoRAAdapter.cs
• src/LoRA/Adapters/DeltaLoRAAdapter.cs
• src/LoRA/Adapters/GLoRAAdapter.cs
• src/LoRA/Adapters/LoRETTAAdapter.cs
• src/LoRA/Adapters/NOLAAdapter.cs
• src/LoRA/Adapters/RoSAAdapter.cs
• src/LoRA/Adapters/TiedLoRAAdapter.cs
• src/LoRA/Adapters/XLoRAAdapter.cs
• src/LoRA/LoRALayer.cs
• src/Models/Results/AiModelResult.cs
• src/NeuralNetworks/Layers/Conv3DLayer.cs
• src/NeuralNetworks/Layers/DecoderLayer.cs
• src/NeuralNetworks/Layers/DenseLayer.cs
• src/NeuralNetworks/Layers/DiffusionConvLayer.cs
• src/NeuralNetworks/Layers/DigitCapsuleLayer.cs
• src/NeuralNetworks/Layers/GRULayer.cs
• src/NeuralNetworks/Layers/GatedLinearUnitLayer.cs
• src/NeuralNetworks/Layers/InteractingLayer.cs
• src/NeuralNetworks/Layers/LocallyConnectedLayer.cs
• src/NeuralNetworks/Layers/PatchEmbeddingLayer.cs
• src/NeuralNetworks/Layers/PrimaryCapsuleLayer.cs
• src/NeuralNetworks/Layers/RBMLayer.cs
• src/NeuralNetworks/Layers/RecurrentLayer.cs
• src/NeuralNetworks/Layers/SSM/GatedLinearAttentionLayer.cs
• src/NeuralNetworks/Layers/SSM/HybridBlockScheduler.cs
• src/NeuralNetworks/Layers/SSM/HyenaLayer.cs
• src/NeuralNetworks/Layers/SSM/MixtureOfMambaLayer.cs
• src/NeuralNetworks/Layers/SSM/S5Layer.cs
• src/NeuralNetworks/Layers/SeparableConvolutionalLayer.cs
• src/NeuralNetworks/Layers/SparseLinearLayer.cs
• src/NeuralNetworks/Layers/SpiralConvLayer.cs
• src/NeuralNetworks/Layers/SubpixelConvolutionalLayer.cs
• src/NeuralNetworks/Layers/TimeEmbeddingLayer.cs
• src/NeuralNetworks/Layers/TimeMoEBlockLayer.cs
• src/NeuralNetworks/NeuralNetworkBase.cs
• src/NeuralNetworks/VideoCLIPNeuralNetwork.cs
• src/OnlineLearning/OnlineLearningModelBase.cs
• src/PhysicsInformed/NeuralOperators/FourierNeuralOperator.cs
• src/PointCloud/Models/PointNetPlusPlus.cs
• src/SelfSupervisedLearning/SymmetricProjector.cs
• src/SurvivalAnalysis/SurvivalModelBase.cs
• tests/AiDotNet.Tests/IntegrationTests/WeightStreaming/WeightStreamingEndToEndTests.cs

[coderabbitai]

Actionable comments posted: 2

♻️ Duplicate comments (1)

src/NeuralNetworks/Layers/MLPMixerBlockLayer.cs (1)

163-204: ⚠️ Potential issue | 🔴 Critical | ⚡ Quick win

Blocking: validate the full parameter vector before mutating sublayers.

Line 191 starts writing into child layers before this method knows parameters.Length matches the composite layout. A short or long vector can therefore leave MLPMixerBlockLayer<T> partially updated before the final throw, which is a bad failure mode for deserialization/load paths.

Proposed fix

 public override void SetParameters(Vector<T> parameters)
 {
     // Lazy ctor: sublayers start with placeholder shapes. Resolve them
     // from known constants (numPatches, hiddenDim, expansion factors).
     int tempExpanded = _numPatches * _expansionFactor;
     int chanExpanded = _hiddenDim * _expansionFactor;

     if (!_norm1.IsShapeResolved) _norm1.ResolveFromShape(new[] { _hiddenDim });
     if (!_norm2.IsShapeResolved) _norm2.ResolveFromShape(new[] { _hiddenDim });
     if (!_temporalMlpExpand.IsShapeResolved) _temporalMlpExpand.ResolveFromShape(new[] { _numPatches });
     if (!_temporalMlpContract.IsShapeResolved) _temporalMlpContract.ResolveFromShape(new[] { tempExpanded });
     if (!_channelMlpExpand.IsShapeResolved) _channelMlpExpand.ResolveFromShape(new[] { _hiddenDim });
     if (!_channelMlpContract.IsShapeResolved) _channelMlpContract.ResolveFromShape(new[] { chanExpanded });
+
+    long expectedTotal =
+        _norm1.ParameterCount +
+        _temporalMlpExpand.ParameterCount +
+        _temporalMlpContract.ParameterCount +
+        _norm2.ParameterCount +
+        _channelMlpExpand.ParameterCount +
+        _channelMlpContract.ParameterCount;
+    if (parameters.Length != expectedTotal)
+    {
+        throw new ArgumentException(
+            $"MLPMixerBlockLayer expected {expectedTotal} parameters across sublayers, got {parameters.Length}.",
+            nameof(parameters));
+    }

     int idx = 0;
     void Set(ILayer<T> sub)
     {
         int count = checked((int)sub.ParameterCount);
         if (count == 0) return;
         sub.SetParameters(parameters.Slice(idx, count));
         idx += count;
     }
     Set(_norm1);
     Set(_temporalMlpExpand);
     Set(_temporalMlpContract);
     Set(_norm2);
     Set(_channelMlpExpand);
     Set(_channelMlpContract);
-    if (idx != parameters.Length)
-    {
-        throw new ArgumentException(
-            $"MLPMixerBlockLayer expected {idx} parameters across sublayers, got {parameters.Length}.");
-    }
 }

As per coding guidelines, "Every PR must contain production-ready code" and "missing validation of external inputs" is blocking.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@src/NeuralNetworks/Layers/MLPMixerBlockLayer.cs` around lines 163 - 204, The
SetParameters method mutates child layers before verifying the incoming
parameters.Length matches the composite expected size; first compute the total
expected parameter count by summing each sub.ParameterCount (using the same
checked((int)sub.ParameterCount) logic for _norm1, _temporalMlpExpand,
_temporalMlpContract, _norm2, _channelMlpExpand, _channelMlpContract), validate
parameters.Length == totalCount and throw if not, and only then perform the
existing loop that slices and calls sub.SetParameters (keeping ResolveFromShape
calls as-is); this ensures no partial updates occur on mismatch.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@src/NeuralNetworks/Layers/DecoderLayer.cs`:
- Around line 149-186: The layer currently guards only SetParameters against a
null _feedForward2, leaving other public members (ParameterCount, GetParameters,
GetParameterGradients, ClearGradients, UpdateParameters, UpdateParametersGpu,
ResetState) vulnerable to NullReferenceException; add a single helper (e.g.
RequireResolvedFeedForward2(string caller) that throws InvalidOperationException
with a clear message when _feedForward2 is null) and call it at the start of
each of those methods (and in SetParameters to replace the existing check) so
all public entry points fail fast and consistently when the layer shape hasn’t
been resolved.
- Around line 115-145: Ensure DecoderLayer.Serialize and Deserialize validate
InputSize and fail fast on unresolved/invalid sizes: in Serialize, check
IsShapeResolved and that InputSize > 0 and throw an exception (e.g.,
InvalidOperationException) if the layer is unresolved or InputSize is
non-positive instead of writing a bad value; in Deserialize, after reading
savedInputSize throw an InvalidDataException if savedInputSize <= 0 before any
state mutation, and only call ResolveFromShape or base.Deserialize when
savedInputSize is positive and consistent with InputSize; update error messages
to reference DecoderLayer, InputSize and IsShapeResolved to make failures
actionable.

---

Duplicate comments:
In `@src/NeuralNetworks/Layers/MLPMixerBlockLayer.cs`:
- Around line 163-204: The SetParameters method mutates child layers before
verifying the incoming parameters.Length matches the composite expected size;
first compute the total expected parameter count by summing each
sub.ParameterCount (using the same checked((int)sub.ParameterCount) logic for
_norm1, _temporalMlpExpand, _temporalMlpContract, _norm2, _channelMlpExpand,
_channelMlpContract), validate parameters.Length == totalCount and throw if not,
and only then perform the existing loop that slices and calls sub.SetParameters
(keeping ResolveFromShape calls as-is); this ensures no partial updates occur on
mismatch.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: ASSERTIVE

Plan: Pro

Run ID: 502d5962-f9c1-4c45-a274-52cbc20241c1

📥 Commits

Reviewing files that changed from the base of the PR and between 5fa4c44 and 1fd05c2.

📒 Files selected for processing (3)

• src/NeuralNetworks/Layers/DecoderLayer.cs
• src/NeuralNetworks/Layers/KairosMultiSizePatchLayer.cs
• src/NeuralNetworks/Layers/MLPMixerBlockLayer.cs

[Copilot]

Copilot encountered an error and was unable to review this pull request. You can try again by re-requesting a review.

[Copilot]

Copilot encountered an error and was unable to review this pull request. You can try again by re-requesting a review.

[Copilot]

Copilot encountered an error and was unable to review this pull request. You can try again by re-requesting a review.

[Copilot]

Copilot encountered an error and was unable to review this pull request. You can try again by re-requesting a review.