perf: lazy attention + Dispose-to-pool cascade (parts 2+3 of #1136) - #1140
Merged
Merged
Conversation
Every diffusion model test on CI was OOMing during construction because DiT/MMDiT/UNet noise predictors eagerly allocate ~4 GB of weight tensors in their ctors — `new DenseLayer<T>(...)` calls `TensorAllocator.Rent<T>` before the model has even seen an input. With 255 diffusion models × ~10 tests each running in a shared xunit process, sequential tests stack up and OOM on the 16 GB Windows CI runners, cancelling all three Diffusion ModelFamily shards at the 45-minute wall clock. This part of the fix makes model construction O(1) by threading `InitializationStrategies<T>.Lazy` into every internal Dense/Conv layer built by a noise predictor. Weight tensors stay at shape [0,0] until the first Forward() pass actually needs them. Adds two protected helpers on NoisePredictorBase<T>: - LazyDense(int, int, IActivationFunction?) — lazy dense layer - LazyDenseVec(int, int, IVectorActivationFunction) — lazy dense for vector activations (distinct name avoids ctor-overload ambiguity on activations that implement both interfaces) - LazyConv2D(...) — lazy 2D convolutional layer Converts all 11 noise predictors + DiffusionResBlock: - DiTNoisePredictor (patch/time embeds + final + MLP/attention blocks) - MMDiTNoisePredictor (image/text streams + joint + single blocks) - MMDiTXNoisePredictor - EMMDiTPredictor - FlagDiTPredictor - FluxDoubleStreamPredictor (double + single streams) - AsymmDiTPredictor - SiTPredictor - UViTNoisePredictor (encoder + decoder + skip projections) - UNetNoisePredictor (downsample) - VideoUNetPredictor (time embed + spatial/temporal res blocks + in/out conv) - DiffusionResBlock (norm/conv1 + timeMlp + norm/conv2 + skip conv) Parts 2-5 (MHA/LayerNorm lazy init, Dispose→pool return, test lifecycle, GetParameters→ParameterCount swap) follow in subsequent commits. Refs #1136
…arts 4+5/5 of #1136) Two follow-on pieces to the lazy-init-in-noise-predictors change (PR 1 of the #1136 series). Together with that first commit they should clear the three cancelled Diffusion ModelFamily shards on CI — if not fully, they put the next diagnostics in reach instead of hiding behind OOM. Part 5: swap `GetParameters().Length > 0` for `ParameterCount > 0`. The `Parameters_ShouldBeNonEmpty` tests in DiffusionModelTestBase and NeuralNetworkModelTestBase are just asking "does this model have any learnable parameters?" — a question `ParameterCount` answers in O(1) without materializing the full flattened parameter vector. Calling GetParameters() on a lazily-constructed DiT-XL forces every DenseLayer to eagerly allocate weights (EnsureInitialized cascades down) just to count them — that alone reproduces the ~4 GB OOM even with lazy init in place. ParameterCount propagates through the layer list without triggering lazy materialization, so the existence check stays cheap. Semantically identical assertion, not a weakening — both check that the model has learnable parameters. Just using the right API. Part 4: IAsyncLifetime.DisposeAsync forces a full GC cycle between test methods in both ModelFamily bases. Each test instantiates a fresh production-sized model (VGG16BN, DiT-XL, SDXL, etc.). Without explicit GC pressure, the shared xunit process holds onto previous test's weight tensors until the collector runs on its own schedule. Sequential tests in a shard stack up live weight allocations that exceed the 16 GB Windows runner budget. IAsyncLifetime.DisposeAsync runs after every [Fact] in a class, so a Collect/WaitForPendingFinalizers/Collect cycle there releases the previous test's weight tensors before the next test's model constructs. Combined with Part 1 (lazy construction) and Part 5 (no forced materialization on existence check), the working-set for a single test stays bounded to what that test actually touches. Parts 2 (MHA/LayerNorm lazy init) and 3 (Dispose → pool-return) are deferred to a follow-up PR once we see how much of the OOM is addressed by Parts 1/4/5 alone. Refs #1136
…3/5 of #1136) Second round of memory fixes for the cancelled Diffusion / NeuralNetworks CI shards. Parts 1+4+5 in PR #1137 landed lazy DenseLayer/ConvolutionalLayer init in the noise predictors but didn't move the needle — Diffusion shards still cancelled at 45 min. The diagnosis: each DiT-XL tower still eagerly allocates ~1 GB of attention weights across 28 blocks' worth of Q/K/V projections, because neither `SelfAttentionLayer` nor `MultiHeadAttentionLayer` had a lazy-init path. This PR adds it. ## Part 2 — Lazy init for MHA + SelfAttentionLayer Both layers gain an optional `IInitializationStrategy<T>?` constructor parameter. When the strategy is `InitializationStrategies<T>.Lazy`, Q/K/V/O weight tensors stay at shape `[0,0]` until the first Forward() call: - `SelfAttentionLayer`: straightforward — `EnsureInitialized()` override allocates on first Forward / GetParameters / SetParameters. No sub-layer fields, so no TrainableParameterGenerator conflict. - `MultiHeadAttentionLayer`: trickier. The generator already emits an `EnsureInitialized` override on MHA because `_ropeLayer` and `_alibiLayer` are sub-layer fields needing registration. Our lazy-init logic lives in a separate private helper `EnsureWeightsAllocated()` called from Forward / GetParameters / SetParameters. The generator-emitted EnsureInitialized still runs its sub-layer registration. Both expose lazy-friendly `ParameterCount` so existence checks don't force allocation. `LazyMHA` and `LazySelfAttention` factories added to `NoisePredictorBase<T>` alongside the existing `LazyDense` / `LazyConv2D` helpers. `DiTNoisePredictor.CreateAttentionLayer`, `UViTNoisePredictor`, and `VideoUNetPredictor.CreateSpatial/Temporal/CrossAttention` updated. Skipped `LayerNormalizationLayer` — gamma/beta are `[featureSize]` tensors (~18 KB each), not a meaningful OOM contributor. ## Part 3 — Dispose cascades rented tensors back to the pool - `ConvolutionalLayer.Dispose` now calls `TensorAllocator.Return(_kernels)` in the eager-init branch, matching the pattern `DenseLayer.Dispose` already has for `_weights` / `_biases`. The convolutional kernels are rented via `RentUninitialized` in the eager path, so they need to go back to the pool for reuse by the next model. - `NeuralNetworkBase.Dispose(bool)` now cascades to every layer that implements `IDisposable`, catching `ObjectDisposedException` so a cross-network layer share doesn't abort the loop. - `DiffusionModelBase` gains its own `Dispose` / `Dispose(bool)` since it doesn't inherit from `NeuralNetworkBase`. Subclasses can override to cascade to their own disposables. - `INeuralNetworkModel<T>` and `IDiffusionModel<T>` now inherit `IDisposable` so callers can wrap models in `using var`. `NeuralNetworkBase` and `DiffusionModelBase` already implement the contract — this just tightens the interface surface to match what concrete types already do. - Test bases `DiffusionModelTestBase` and `NeuralNetworkModelTestBase` now wrap `CreateModel()` / `CreateNetwork()` results in `using var` for all 29 test methods. Combined with the `IAsyncLifetime.DisposeAsync` GC hook from PR #1137, each test now actually releases its weight buffers back to the allocator pool before the next test constructs. ## Stacked on #1137 Branched off `perf/diffusion-lazy-init-oom`. Merges naturally either way (this PR's diff against master after #1137 lands cleanly covers only the Parts 2+3 additions). Refs #1136
|
The latest updates on your projects. Learn more about Vercel for GitHub. 2 Skipped Deployments
|
Contributor
WalkthroughThis PR introduces lazy weight initialization across core layers (SelfAttention, MultiHeadAttention, Dense, Conv), adds lazy layer factory helpers for noise predictors, updates several noise predictors to use lazy factories, strengthens disposal semantics (INeuralNetworkModel → IDisposable and cascade disposal), and adjusts tests to avoid parameter materialization. Changes
Sequence DiagramsequenceDiagram
participant Model as Predictor / Model
participant Layer as Lazy Layer (Self/MHA/Dense/Conv)
participant Strategy as InitializationStrategy
participant Alloc as TensorAllocator
Model->>Layer: Construct(..., initializationStrategy=Lazy)
Layer->>Layer: Allocate zero-sized placeholder tensors\n_set _isInitialized = false_
Note over Layer: No heavy memory allocation at construction
Model->>Layer: Forward(input)
Layer->>Layer: EnsureInitialized() / EnsureWeightsAllocated()
alt Weights not materialized
Layer->>Alloc: Request real tensor shapes
Alloc-->>Layer: Provide tensors
Layer->>Strategy: InitializeWeights()/InitializeBiases()
Strategy-->>Layer: Fill tensor values
Layer->>Layer: Register trainable parameters\n_set _isInitialized = true
end
Layer->>Layer: Run computation using materialized tensors
Estimated code review effort🎯 4 (Complex) | ⏱️ ~50 minutes Possibly related issues
Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Comment |
This was referenced Apr 14, 2026
ooples
added a commit
that referenced
this pull request
Apr 14, 2026
EmbeddingLayer's tensor was eagerly allocated in the constructor — [vocabularySize, embeddingDimension]. For BERT-scale transformers (BGE / SGPT / Matryoshka), that's ~30,522 × 768 × 8 bytes ≈ 187 MB materialized at `new EmbeddingLayer(...)` time, before any test input flows through. Multiple EmbeddingLayer instances per network × multiple test methods × shared xunit process = hundreds of MB held live across test boundaries. Made initialization lazy following the same pattern PR #1140 introduced for MultiHeadAttentionLayer / SelfAttentionLayer: - Cached _vocabularySize / _embeddingDimension / _embeddingInitialized fields. ParameterCount reads from the cached fields, not from the placeholder tensor's [0,0] shape — stays correct without forcing materialization. - Constructor leaves _embeddingTensor at [0,0] until first access. EnsureEmbeddingInitialized() allocates the real shape, runs the same SimdRandom-scaled fill the constructor used to do, registers with the engine for GPU persistence. - Wired into Forward, GetParameters, GetTokenEmbeddings, SetParameters so any data-touching path materializes first. SetParameters also short-circuits the lazy flag since it's writing the real-sized tensor itself. Refs #1136
ooples
added a commit
that referenced
this pull request
Apr 15, 2026
…osableComponents Closes the diffusion-side Dispose gap. Previously DiffusionModelBase implemented neither IDisposable nor a cascade — concrete diffusion models (249 of them) couldn't release their composed predictor's pool-rented weight tensors back to the allocator without each implementing Dispose by hand. This caused the test-shard OOMs in the earlier perf work even after the layer-side Dispose plumbing (commit a9070a6 in PR #1140). DiffusionModelBase changes: - Implements IDisposable. No existing concrete diffusion subclass declared its own Dispose, so this is purely additive. - protected virtual EnumerateDisposableComponents() — concrete models override to yield the components they own (predictor, VAE, conditioner). Default returns empty so the 249 existing subclasses keep working unchanged until each opts in. Migration is per-model: protected override IEnumerable<IDisposable> EnumerateDisposableComponents() { yield return _unet; yield return _vae; } - Dispose(bool) walks the enumeration, catching ObjectDisposedException for shared-component graphs (a predictor reused across two diffusion wrappers for ensembling). Why a method-based opt-in, not a NoisePredictor property: ILatentDiffusionModel<T> already declares INoisePredictor<T> NoisePredictor { get; } as a non-nullable contract that LatentDiffusionModelBase implements. Adding a virtual nullable property on the base would either collide on the name (CS0114), require nullability gymnastics across every latent-diffusion subclass, or break the existing interface contract. The method-based pattern sidesteps the collision entirely and matches the EnumerateLayers hook that NoisePredictorBase introduced in the prior commit. Build clean on net10.0 + net471, 0 errors. Behavior unchanged for all existing concrete diffusion models. Part B.3 of the lucky-waddling-knuth plan. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
ooples
added a commit
that referenced
this pull request
Apr 15, 2026
…y-allocs, internal shape access Critical: - DiTNoisePredictor constructor now throws NotSupportedException when numClasses > 0. The _labelEmbed layer is created and participates in param/grad/serialize plumbing but no forward path actually injects the class index — weights would train against zero-gradient signal silently. Every caller in the codebase currently passes numClasses=0 (the feature has zero active users), so the throw is safe. The plumbing stays so a future PR can wire the class-injection path end-to-end (following DiT paper §3.2: fuse label embedding into projected time embedding before AdaLN modulation). Major: - VideoUNetPredictor._inputConv, _outputConv, _timeEmbedMlp1, _timeEmbedMlp2, _imageCondProjection, CreateSpatialResBlock, and CreateDownsample now use the Lazy factories. Previously only 2 of ~10 layer factories were lazy, re-introducing the exact OOM pressure the PR was meant to relieve. MHA attention layers and DeconvolutionalLayer don't have Lazy factories on this branch yet (attention lazy support is on PR #1140; no LazyDeconv exists), so those remain eager — documented as a follow-up. - ApplyFiLMConditioning validates projection output channels match x's channel dim. Without this check, a misconfigured projection silently produces wrong-shape broadcasts or broadcasts into the wrong axis. - ApplyTemporalProcessing and ApplyTemporalAttention now use video.Shape (public) instead of video._shape (internal backing field). Decouples the predictor from Tensor<T>'s storage layout. Build clean on net10.0 + net471, 0 errors. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
ooples
added a commit
that referenced
this pull request
Apr 15, 2026
…erate-disposable-components Closes the diffusion-side Dispose gap. Previously DiffusionModelBase implemented neither IDisposable nor a cascade — concrete diffusion models (249 of them) couldn't release their composed predictor's pool-rented weight tensors back to the allocator without each implementing Dispose by hand. This caused the test-shard OOMs in the earlier perf work even after the layer-side Dispose plumbing (commit a9070a6 in PR #1140). DiffusionModelBase changes: - Implements IDisposable. No existing concrete diffusion subclass declared its own Dispose, so this is purely additive. - protected virtual EnumerateDisposableComponents() — concrete models override to yield the components they own (predictor, VAE, conditioner). Default returns empty so the 249 existing subclasses keep working unchanged until each opts in. Migration is per-model: protected override IEnumerable<IDisposable> EnumerateDisposableComponents() { yield return _unet; yield return _vae; } - Dispose(bool) walks the enumeration, catching ObjectDisposedException for shared-component graphs (a predictor reused across two diffusion wrappers for ensembling). Why a method-based opt-in, not a NoisePredictor property: ILatentDiffusionModel<T> already declares INoisePredictor<T> NoisePredictor { get; } as a non-nullable contract that LatentDiffusionModelBase implements. Adding a virtual nullable property on the base would either collide on the name (CS0114), require nullability gymnastics across every latent-diffusion subclass, or break the existing interface contract. The method-based pattern sidesteps the collision entirely and matches the EnumerateLayers hook that NoisePredictorBase introduced in the prior commit. Build clean on net10.0 + net471, 0 errors. Behavior unchanged for all existing concrete diffusion models. Part B.3 of the lucky-waddling-knuth plan. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
ooples
added a commit
that referenced
this pull request
Apr 15, 2026
…y-allocs, internal shape access Critical: - DiTNoisePredictor constructor now throws NotSupportedException when numClasses > 0. The _labelEmbed layer is created and participates in param/grad/serialize plumbing but no forward path actually injects the class index — weights would train against zero-gradient signal silently. Every caller in the codebase currently passes numClasses=0 (the feature has zero active users), so the throw is safe. The plumbing stays so a future PR can wire the class-injection path end-to-end (following DiT paper §3.2: fuse label embedding into projected time embedding before AdaLN modulation). Major: - VideoUNetPredictor._inputConv, _outputConv, _timeEmbedMlp1, _timeEmbedMlp2, _imageCondProjection, CreateSpatialResBlock, and CreateDownsample now use the Lazy factories. Previously only 2 of ~10 layer factories were lazy, re-introducing the exact OOM pressure the PR was meant to relieve. MHA attention layers and DeconvolutionalLayer don't have Lazy factories on this branch yet (attention lazy support is on PR #1140; no LazyDeconv exists), so those remain eager — documented as a follow-up. - ApplyFiLMConditioning validates projection output channels match x's channel dim. Without this check, a misconfigured projection silently produces wrong-shape broadcasts or broadcasts into the wrong axis. - ApplyTemporalProcessing and ApplyTemporalAttention now use video.Shape (public) instead of video._shape (internal backing field). Decouples the predictor from Tensor<T>'s storage layout. Build clean on net10.0 + net471, 0 errors. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
1 of 3 tasks
ooples
added a commit
that referenced
this pull request
Apr 16, 2026
…ffusion (#1143) * refactor: extract compile cache into composable CompiledModelHost<T> One attachment point for every future compilation feature — AOT plan serialization, CUDA Graph capture, symbolic shape plans, persistent autotune — rather than re-implementing the compile+cache+fallback dance per model family. New src/NeuralNetworks/CompiledModelHost.cs: - Sealed component, not a base class. Consumers own the layer list and Predict API; they compose one host instance. - Predict(input, structureVersion, eagerForward) — traces on miss, replays on hit, falls back to eager on disable/failure. - Automatic invalidation on structureVersion mismatch — callers bump the version when their layer graph mutates (lazy-init resize, layer add/remove, weight swap) and stale plans get dropped before the next compile. Previously the per-model code had to remember to call _compiledInferenceCache.Invalidate() on every mutation path; now it's centralized. NeuralNetworkBase: - Replaces the ad-hoc `_compiledInferenceCache` field and PredictCompiled body with `_compileHost.Predict(input, _layerStructureVersion, ...)`. PredictCompiled becomes a one-line delegation. - Dispose(bool) now cascades to every layer implementing IDisposable so pool-rented weight tensors return to the allocator before GC. Wrapped in try/catch(ObjectDisposedException) so a shared-layer graph doesn't abort the cascade when a previous owner already disposed the layer. - Compile host disposed first, then layers, then DisableMixedPrecision. Build clean on net10.0 + net471, 0 errors. This is the foundation commit. Follow-ups on the same branch or new PRs: - Wire NoisePredictorBase to the same host (closes diffusion compile gap) - Migrate 11 concrete predictor subclasses to register layers with compile-friendly hooks - DiffusionModelBase exposes NoisePredictor property so AiModelBuilder's BuildCompiledPredictFunction can reach through Part B.1 of the lucky-waddling-knuth plan. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * refactor: extract compile cache into composable CompiledModelHost<T> One attachment point for every future compilation feature — AOT plan serialization, CUDA Graph capture, symbolic shape plans, persistent autotune — rather than re-implementing the compile+cache+fallback dance per model family. New src/NeuralNetworks/CompiledModelHost.cs: - Sealed component, not a base class. Consumers own the layer list and Predict API; they compose one host instance. - Predict(input, structureVersion, eagerForward) — traces on miss, replays on hit, falls back to eager on disable/failure. - Automatic invalidation on structureVersion mismatch — callers bump the version when their layer graph mutates (lazy-init resize, layer add/remove, weight swap) and stale plans get dropped before the next compile. Previously the per-model code had to remember to call _compiledInferenceCache.Invalidate() on every mutation path; now it's centralized. NeuralNetworkBase: - Replaces the ad-hoc `_compiledInferenceCache` field and PredictCompiled body with `_compileHost.Predict(input, _layerStructureVersion, ...)`. PredictCompiled becomes a one-line delegation. - Dispose(bool) now cascades to every layer implementing IDisposable so pool-rented weight tensors return to the allocator before GC. Wrapped in try/catch(ObjectDisposedException) so a shared-layer graph doesn't abort the cascade when a previous owner already disposed the layer. - Compile host disposed first, then layers, then DisableMixedPrecision. Build clean on net10.0 + net471, 0 errors. This is the foundation commit. Follow-ups on the same branch or new PRs: - Wire NoisePredictorBase to the same host (closes diffusion compile gap) - Migrate 11 concrete predictor subclasses to register layers with compile-friendly hooks - DiffusionModelBase exposes NoisePredictor property so AiModelBuilder's BuildCompiledPredictFunction can reach through Part B.1 of the lucky-waddling-knuth plan. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * refactor: wire NoisePredictorBase through CompiledModelHost + Dispose Closes the diffusion-side half of the compile-plan gap. Previously DiffusionModelBase.NoisePredictor could not benefit from compiled-plan replay because the predictor base had no compile cache. Now every concrete predictor (DiT, MMDiT, UNet variants — 11 subclasses) can opt into compiled inference for its per-step Forward by calling the new protected PredictCompiled helper, and the 50-step denoising loop gets a near-zero-overhead replay after one trace instead of eager full-graph execution on every step. NoisePredictorBase changes: - Implements IDisposable. No existing predictor subclass declared its own Dispose, so this is a pure addition to the interface surface. - Owns a private CompiledModelHost<T> (the sealed component landed in the prior commit). One attachment point; everything else — AOT plan serialization, CUDA Graph capture, persistent autotune — attaches here too. - _layerStructureVersion + InvalidateCompiledPlans() hook for concrete predictors to signal "the layer graph just changed, drop stale plans." Called from lazy-init paths and SetParameters swaps. - protected virtual EnumerateLayers() — concrete predictors override to expose their private layer fields (_blocks, _timeEmbed1, _patchEmbed, etc.) for Dispose cascade. Default returns empty so this commit is non-breaking; migration of the 11 concrete predictors to populate EnumerateLayers lands in follow-up commits on this same branch. - PredictCompiled(input, eagerFallback) — one-line delegation to the host. Concrete predictors route their Forward through this once they're verified safe for trace+replay. - Dispose cascade: host first (so pool-captured tensors are freed), then every ILayer<T> from EnumerateLayers that implements IDisposable. Wrapped in try/catch(ObjectDisposedException) for shared-layer graphs. Build clean on net10.0 + net471, 0 errors. Behavior unchanged for all existing callers — the opt-in hooks are additive. Part B.2 of the lucky-waddling-knuth plan. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * refactor: wire NoisePredictorBase through CompiledModelHost + Dispose Closes the diffusion-side half of the compile-plan gap. Previously DiffusionModelBase.NoisePredictor could not benefit from compiled-plan replay because the predictor base had no compile cache. Now every concrete predictor (DiT, MMDiT, UNet variants — 11 subclasses) can opt into compiled inference for its per-step Forward by calling the new protected PredictCompiled helper, and the 50-step denoising loop gets a near-zero-overhead replay after one trace instead of eager full-graph execution on every step. NoisePredictorBase changes: - Implements IDisposable. No existing predictor subclass declared its own Dispose, so this is a pure addition to the interface surface. - Owns a private CompiledModelHost<T> (the sealed component landed in the prior commit). One attachment point; everything else — AOT plan serialization, CUDA Graph capture, persistent autotune — attaches here too. - _layerStructureVersion + InvalidateCompiledPlans() hook for concrete predictors to signal "the layer graph just changed, drop stale plans." Called from lazy-init paths and SetParameters swaps. - protected virtual EnumerateLayers() — concrete predictors override to expose their private layer fields (_blocks, _timeEmbed1, _patchEmbed, etc.) for Dispose cascade. Default returns empty so this commit is non-breaking; migration of the 11 concrete predictors to populate EnumerateLayers lands in follow-up commits on this same branch. - PredictCompiled(input, eagerFallback) — one-line delegation to the host. Concrete predictors route their Forward through this once they're verified safe for trace+replay. - Dispose cascade: host first (so pool-captured tensors are freed), then every ILayer<T> from EnumerateLayers that implements IDisposable. Wrapped in try/catch(ObjectDisposedException) for shared-layer graphs. Build clean on net10.0 + net471, 0 errors. Behavior unchanged for all existing callers — the opt-in hooks are additive. Part B.2 of the lucky-waddling-knuth plan. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * refactor: add opt-in dispose cascade to diffusion model base via enumerate-disposable-components Closes the diffusion-side Dispose gap. Previously DiffusionModelBase implemented neither IDisposable nor a cascade — concrete diffusion models (249 of them) couldn't release their composed predictor's pool-rented weight tensors back to the allocator without each implementing Dispose by hand. This caused the test-shard OOMs in the earlier perf work even after the layer-side Dispose plumbing (commit a9070a6 in PR #1140). DiffusionModelBase changes: - Implements IDisposable. No existing concrete diffusion subclass declared its own Dispose, so this is purely additive. - protected virtual EnumerateDisposableComponents() — concrete models override to yield the components they own (predictor, VAE, conditioner). Default returns empty so the 249 existing subclasses keep working unchanged until each opts in. Migration is per-model: protected override IEnumerable<IDisposable> EnumerateDisposableComponents() { yield return _unet; yield return _vae; } - Dispose(bool) walks the enumeration, catching ObjectDisposedException for shared-component graphs (a predictor reused across two diffusion wrappers for ensembling). Why a method-based opt-in, not a NoisePredictor property: ILatentDiffusionModel<T> already declares INoisePredictor<T> NoisePredictor { get; } as a non-nullable contract that LatentDiffusionModelBase implements. Adding a virtual nullable property on the base would either collide on the name (CS0114), require nullability gymnastics across every latent-diffusion subclass, or break the existing interface contract. The method-based pattern sidesteps the collision entirely and matches the EnumerateLayers hook that NoisePredictorBase introduced in the prior commit. Build clean on net10.0 + net471, 0 errors. Behavior unchanged for all existing concrete diffusion models. Part B.3 of the lucky-waddling-knuth plan. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * refactor: DiffusionModelBase opt-in Dispose cascade via EnumerateDisposableComponents Closes the diffusion-side Dispose gap. Previously DiffusionModelBase implemented neither IDisposable nor a cascade — concrete diffusion models (249 of them) couldn't release their composed predictor's pool-rented weight tensors back to the allocator without each implementing Dispose by hand. This caused the test-shard OOMs in the earlier perf work even after the layer-side Dispose plumbing (commit a9070a6 in PR #1140). DiffusionModelBase changes: - Implements IDisposable. No existing concrete diffusion subclass declared its own Dispose, so this is purely additive. - protected virtual EnumerateDisposableComponents() — concrete models override to yield the components they own (predictor, VAE, conditioner). Default returns empty so the 249 existing subclasses keep working unchanged until each opts in. Migration is per-model: protected override IEnumerable<IDisposable> EnumerateDisposableComponents() { yield return _unet; yield return _vae; } - Dispose(bool) walks the enumeration, catching ObjectDisposedException for shared-component graphs (a predictor reused across two diffusion wrappers for ensembling). Why a method-based opt-in, not a NoisePredictor property: ILatentDiffusionModel<T> already declares INoisePredictor<T> NoisePredictor { get; } as a non-nullable contract that LatentDiffusionModelBase implements. Adding a virtual nullable property on the base would either collide on the name (CS0114), require nullability gymnastics across every latent-diffusion subclass, or break the existing interface contract. The method-based pattern sidesteps the collision entirely and matches the EnumerateLayers hook that NoisePredictorBase introduced in the prior commit. Build clean on net10.0 + net471, 0 errors. Behavior unchanged for all existing concrete diffusion models. Part B.3 of the lucky-waddling-knuth plan. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * fix: address PR #1143 review feedback — Dispose completeness + observability Critical (blocking): - DiffusionModelBase.EnumerateDisposableComponents default is now a reflection walk over instance fields that yields anything implementing IDisposable. Previously it returned empty so all 249 concrete diffusion subclasses kept their predictors/VAEs/conditioners alive after Dispose. Subclasses can still override for an explicit allow-list (preferred in perf-sensitive code). - NoisePredictorBase.EnumerateLayers default is the analogous reflection walk that finds ILayer<T> fields and ILayer<T> entries inside collection fields (catches _blocks, _timeEmbed1, _patchEmbed, etc.). Concrete predictors no longer leak their network graph by default. - NeuralNetworkBase.Dispose calls DisableMemoryManagement() before DisableMixedPrecision() so activation pool / gradient checkpoint state is freed instead of surviving model disposal. - All ObjectDisposedException catches in the cascade now Trace.TraceWarning with the layer/component name + exception message instead of fully swallowing — double-dispose bugs become diagnosable in production telemetry. Major: - DiffusionModelBase.Dispose now disposes _scheduler when it implements IDisposable. Previously schedulers with native handles or precomputed alpha/beta buffers leaked across model disposal. - NoisePredictorBase.Dispose releases the _timestepEmbeddingCache tensors — these are owned exclusively by the predictor and had no other Dispose path. - CompiledModelHost.Predict catch now invalidates the entire cache after a failure (replay path could leave a broken plan cached, causing every subsequent call to re-enter the same exception). Resets the version watermark so the next call recompiles fresh. Trace.TraceWarning logs the failure for observability. - CompiledModelHost is now `internal sealed` instead of `public sealed` — it's plumbing behind the model base classes, not a user-consumed type. Reflection walks use TensorReferenceComparer<object>.Instance for identity tracking (avoids double-yielding the same disposable reachable from two fields, e.g., a predictor stored as both interface and concrete alias). Works on net471 (no ReferenceEqualityComparer there). Skips primitives, strings, value types. Build clean on net10.0 + net471, 0 errors. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * fix: address PR #1143 review feedback — Dispose completeness + observability Critical (blocking): - DiffusionModelBase.EnumerateDisposableComponents default is now a reflection walk over instance fields that yields anything implementing IDisposable. Previously it returned empty so all 249 concrete diffusion subclasses kept their predictors/VAEs/conditioners alive after Dispose. Subclasses can still override for an explicit allow-list (preferred in perf-sensitive code). - NoisePredictorBase.EnumerateLayers default is the analogous reflection walk that finds ILayer<T> fields and ILayer<T> entries inside collection fields (catches _blocks, _timeEmbed1, _patchEmbed, etc.). Concrete predictors no longer leak their network graph by default. - NeuralNetworkBase.Dispose calls DisableMemoryManagement() before DisableMixedPrecision() so activation pool / gradient checkpoint state is freed instead of surviving model disposal. - All ObjectDisposedException catches in the cascade now Trace.TraceWarning with the layer/component name + exception message instead of fully swallowing — double-dispose bugs become diagnosable in production telemetry. Major: - DiffusionModelBase.Dispose now disposes _scheduler when it implements IDisposable. Previously schedulers with native handles or precomputed alpha/beta buffers leaked across model disposal. - NoisePredictorBase.Dispose releases the _timestepEmbeddingCache tensors — these are owned exclusively by the predictor and had no other Dispose path. - CompiledModelHost.Predict catch now invalidates the entire cache after a failure (replay path could leave a broken plan cached, causing every subsequent call to re-enter the same exception). Resets the version watermark so the next call recompiles fresh. Trace.TraceWarning logs the failure for observability. - CompiledModelHost is now `internal sealed` instead of `public sealed` — it's plumbing behind the model base classes, not a user-consumed type. Reflection walks use TensorReferenceComparer<object>.Instance for identity tracking (avoids double-yielding the same disposable reachable from two fields, e.g., a predictor stored as both interface and concrete alias). Works on net471 (no ReferenceEqualityComparer there). Skips primitives, strings, value types. Build clean on net10.0 + net471, 0 errors. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * fix: address PR #1143 round-2 review feedback — Dispose hardening + cache concurrency Major: - DiffusionModelBase.ReflectInstanceDisposables skips _scheduler (handled explicitly by Dispose) so it isn't double-disposed via the cascade. - DiffusionModelBase + NoisePredictorBase reflection helpers Trace.TraceWarning the field name + exception when GetValue throws, instead of silently skipping. Without the trace a private field whose getter throws would leak its disposable resource with no diagnostic at Dispose time. - NoisePredictorBase.InvalidateCompiledPlans now calls _compileHost.Invalidate() immediately rather than waiting for the next PredictCompiled to detect the version mismatch. Releases captured tensor buffers eagerly when the caller's intent is "drop everything, reclaim memory now." - CompiledModelHost adds _sync object and protects all _cache / _lastCompiledVersion / _disposed mutations under lock. Predict acquires the lock briefly to take the cache snapshot, then releases before the potentially-slow compile/replay (CompiledModelCache itself is thread-safe for distinct shapes). Concurrent Dispose/Invalidate/Predict on the same model instance (model serving in a request pool) no longer tears down the cache mid-call. - NeuralNetworkBase.SetParameters now invalidates _compileHost after redistributing parameters. Some ITrainableLayer implementations swap parameter tensors wholesale instead of mutating in place — without invalidation, compiled plans replay against captured tensor references that no longer hold the new weights. Build clean on net10.0 + net471, 0 errors. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * fix: address PR #1143 round-2 review feedback — Dispose hardening + cache concurrency Major: - DiffusionModelBase.ReflectInstanceDisposables skips _scheduler (handled explicitly by Dispose) so it isn't double-disposed via the cascade. - DiffusionModelBase + NoisePredictorBase reflection helpers Trace.TraceWarning the field name + exception when GetValue throws, instead of silently skipping. Without the trace a private field whose getter throws would leak its disposable resource with no diagnostic at Dispose time. - NoisePredictorBase.InvalidateCompiledPlans now calls _compileHost.Invalidate() immediately rather than waiting for the next PredictCompiled to detect the version mismatch. Releases captured tensor buffers eagerly when the caller's intent is "drop everything, reclaim memory now." - CompiledModelHost adds _sync object and protects all _cache / _lastCompiledVersion / _disposed mutations under lock. Predict acquires the lock briefly to take the cache snapshot, then releases before the potentially-slow compile/replay (CompiledModelCache itself is thread-safe for distinct shapes). Concurrent Dispose/Invalidate/Predict on the same model instance (model serving in a request pool) no longer tears down the cache mid-call. - NeuralNetworkBase.SetParameters now invalidates _compileHost after redistributing parameters. Some ITrainableLayer implementations swap parameter tensors wholesale instead of mutating in place — without invalidation, compiled plans replay against captured tensor references that no longer hold the new weights. Build clean on net10.0 + net471, 0 errors. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * fix: address PR #1143 round-3 — eager fallback out of lock + tape-cache invalidation - CompiledModelHost.Predict moves the eager fallback OUT of the _sync lock. Previously _disposed/!EnableCompilation branches called eagerForward() while holding _sync — blocking other threads on what may be a slow op. Now: lock decides which path to take, releases, and the caller-supplied delegate runs unblocked. - NeuralNetworkBase.SetParameters now invalidates BOTH the compile host AND TapeTrainingStep<T>'s collected-parameter cache. The tape cache also keys off captured tensor references — when a layer swaps parameter tensors wholesale, the next training step would replay against parameters that no longer exist. Build clean on net10.0 + net471, 0 errors. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * fix: address PR #1143 round-3 — eager fallback out of lock + tape-cache invalidation - CompiledModelHost.Predict moves the eager fallback OUT of the _sync lock. Previously _disposed/!EnableCompilation branches called eagerForward() while holding _sync — blocking other threads on what may be a slow op. Now: lock decides which path to take, releases, and the caller-supplied delegate runs unblocked. - NeuralNetworkBase.SetParameters now invalidates BOTH the compile host AND TapeTrainingStep<T>'s collected-parameter cache. The tape cache also keys off captured tensor references — when a layer swaps parameter tensors wholesale, the next training step would replay against parameters that no longer exist. Build clean on net10.0 + net471, 0 errors. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * fix: address PR #1143 round-4 — IEnumerable traversal + docs match reflection-walk default - ReflectInstanceDisposables now traverses IEnumerable fields for nested disposables (List<IDisposable>, Dictionary<K, IDisposable>, etc.), matching ReflectInstanceLayers' behavior. Without this, a predictor holding its sub-models in a List<IDisposable> would leak them even though the field itself implements IEnumerable. - Updated XML docs on DiffusionModelBase.Dispose and NoisePredictorBase.Dispose to accurately describe the current behavior: reflection-walk is the DEFAULT, subclasses override to CONSTRAIN (return an explicit allow-list) — the previous docs still said 'subclasses must override to opt in,' which no longer matches reality. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * fix: address PR #1143 round-4 — IEnumerable traversal + docs match reflection-walk default - ReflectInstanceDisposables now traverses IEnumerable fields for nested disposables (List<IDisposable>, Dictionary<K, IDisposable>, etc.), matching ReflectInstanceLayers' behavior. Without this, a predictor holding its sub-models in a List<IDisposable> would leak them even though the field itself implements IEnumerable. - Updated XML docs on DiffusionModelBase.Dispose and NoisePredictorBase.Dispose to accurately describe the current behavior: reflection-walk is the DEFAULT, subclasses override to CONSTRAIN (return an explicit allow-list) — the previous docs still said 'subclasses must override to opt in,' which no longer matches reality. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * fix: address PR #1143 round-3 — DisposeOnceGuard + trace lambda + reflector doc Five reviewer-flagged concerns on the composable compile-host foundation: 1. CompiledModelHost.Predict trace lambda `() => { eagerForward(); }` discarded the Tensor<T> return, which can select a different GetOrCompileInference overload and leave the output tensor ambiguous to the tracer. Changed to `() => eagerForward()` so the output threads through. 2. Introduce DisposeOnceGuard (new Helpers/DisposeOnceGuard.cs) — a process-wide ConditionalWeakTable-backed registry that guarantees each IDisposable instance is disposed at most once, regardless of how many owners cascade into it. Replaces the previous ObjectDisposedException-as-signal pattern, which was unsafe because many layers (e.g., DenseLayer returning pooled tensors to TensorAllocator) are not idempotent on second Dispose. 3. Wire DisposeOnceGuard through: - NeuralNetworkBase.Dispose layer cascade - DiffusionModelBase.Dispose scheduler + component cascade - NoisePredictorBase.Dispose layer cascade Shared-layer graphs (same layer in multiple networks, predictor reused across diffusion wrappers, VAE injected into several models) are now safe. 4. Clarify NoisePredictorBase.EnumerateLayers doc: the reflector walks fields + IEnumerable elements but does NOT recurse into container types whose elements aren't themselves ILayer<T>. Predictors storing layers in container types like `List<DiTBlock>` (where DiTBlock holds layer properties) must override EnumerateLayers. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * fix: address PR #1143 round-3 — DisposeOnceGuard + trace lambda + reflector doc Five reviewer-flagged concerns on the composable compile-host foundation: 1. CompiledModelHost.Predict trace lambda `() => { eagerForward(); }` discarded the Tensor<T> return, which can select a different GetOrCompileInference overload and leave the output tensor ambiguous to the tracer. Changed to `() => eagerForward()` so the output threads through. 2. Introduce DisposeOnceGuard (new Helpers/DisposeOnceGuard.cs) — a process-wide ConditionalWeakTable-backed registry that guarantees each IDisposable instance is disposed at most once, regardless of how many owners cascade into it. Replaces the previous ObjectDisposedException-as-signal pattern, which was unsafe because many layers (e.g., DenseLayer returning pooled tensors to TensorAllocator) are not idempotent on second Dispose. 3. Wire DisposeOnceGuard through: - NeuralNetworkBase.Dispose layer cascade - DiffusionModelBase.Dispose scheduler + component cascade - NoisePredictorBase.Dispose layer cascade Shared-layer graphs (same layer in multiple networks, predictor reused across diffusion wrappers, VAE injected into several models) are now safe. 4. Clarify NoisePredictorBase.EnumerateLayers doc: the reflector walks fields + IEnumerable elements but does NOT recurse into container types whose elements aren't themselves ILayer<T>. Predictors storing layers in container types like `List<DiTBlock>` (where DiTBlock holds layer properties) must override EnumerateLayers. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * fix: address PR #1143 round-4 — dictionary walk + atomic cache swap + narrow catch Three reviewer concerns flagged by CodeRabbit on the round-3 changes: 1. Dispose reflection walks in DiffusionModelBase and NoisePredictorBase handled IEnumerable generically but Dictionary<K, V>.GetEnumerator yields KeyValuePair<K,V>, not values — so disposables/layers stored as dictionary values were silently skipped. Add an explicit IDictionary branch (via DictionaryEntry) before the IEnumerable branch in both walks. 2. CompiledModelHost previously invalidated its CompiledModelCache<T> in place on structureVersion changes, but Predict releases _sync before GetOrCompileInference, so an older in-flight call could repopulate the same invalidated cache after a newer call had already bumped the version. Replace the cache instance instead — Dispose the old cache and allocate a fresh one while still holding _sync. The older call now holds a detached reference with no cross-version bleed. 3. Narrow CompiledModelHost.Predict's fallback catch with an exception filter so fatal CLR failures (OOM, AccessViolation, StackOverflow, BadImageFormat, InvalidProgram, ThreadAbort, AppDomainUnloaded, CannotUnloadAppDomain) propagate instead of being masked by the eager fallback. Allocating again under a poisoned process state would just crash a second way. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * fix: address PR #1143 round-4 — dictionary walk + atomic cache swap + narrow catch Three reviewer concerns flagged by CodeRabbit on the round-3 changes: 1. Dispose reflection walks in DiffusionModelBase and NoisePredictorBase handled IEnumerable generically but Dictionary<K, V>.GetEnumerator yields KeyValuePair<K,V>, not values — so disposables/layers stored as dictionary values were silently skipped. Add an explicit IDictionary branch (via DictionaryEntry) before the IEnumerable branch in both walks. 2. CompiledModelHost previously invalidated its CompiledModelCache<T> in place on structureVersion changes, but Predict releases _sync before GetOrCompileInference, so an older in-flight call could repopulate the same invalidated cache after a newer call had already bumped the version. Replace the cache instance instead — Dispose the old cache and allocate a fresh one while still holding _sync. The older call now holds a detached reference with no cross-version bleed. 3. Narrow CompiledModelHost.Predict's fallback catch with an exception filter so fatal CLR failures (OOM, AccessViolation, StackOverflow, BadImageFormat, InvalidProgram, ThreadAbort, AppDomainUnloaded, CannotUnloadAppDomain) propagate instead of being masked by the eager fallback. Allocating again under a poisoned process state would just crash a second way. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * fix: address pr #1143 review round-3 — dict-cascade + per-version cache + narrow catch + tests Addresses 5 unresolved PR #1143 comments: 1. DiffusionModelBase.ReflectInstanceDisposables: dict-held disposables were silently leaked. Dictionary<TKey, IDisposable> yields KeyValuePair<,> structs (not IDisposable) when enumerated, so the generic IEnumerable branch never reached the payload. Added an IDictionary branch BEFORE the IEnumerable branch that walks DictionaryEntry and extracts .Value, so dictionary-keyed disposables cascade through Dispose. 2. NoisePredictorBase.ReflectInstanceLayers: identical bug for layer enumeration. Dictionary<TKey, ILayer<T>> was silently dropping layers because the enumerator yields KeyValuePair<,> structs. Same fix — IDictionary branch before IEnumerable branch, extracting Value and matching it against ILayer<T>. 3. CompiledModelHost.Predict (structureVersion invalidation): previous code invalidated the cache IN PLACE when structureVersion changed, then reused the same CompiledModelCache<T> instance. A racing Predict that had just released the lock with an older version could repopulate the same cache after the switch, leaking version-N plans into version-N+1 callers. Now we replace the cache reference with a fresh instance per structure version — every per-version local `cache` closes over a unique object, so once we leave the lock it's frozen to that version. 4. CompiledModelHost.Predict (narrow catch): the fallback-to-eager path previously used catch-all `catch (Exception)`, which would swallow fatal CLR exceptions (OutOfMemoryException, StackOverflowException, AccessViolationException, thread aborts) and then immediately invoke eagerForward() on a poisoned process. Replaced with an exception filter that propagates fatal exceptions and only recovers from compile/replay failures. Follows the same exclusion list the runtime team uses for "recoverable vs. non-recoverable" distinctions. 5. CompiledModelHostTests: 9 new unit tests covering the host's invalidation contract, Dispose idempotency, eager fallback paths (compilation disabled, post-Dispose), structureVersion-bump cache invalidation, explicit Invalidate(), null-guard on eagerForward, concurrent Predict+Invalidate, and concurrent Predict+Dispose. All tests passing locally on net10.0. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * fix: PR #1143 review round-3 — dict-cascade + per-version cache + narrow catch + tests Addresses 5 unresolved PR #1143 comments: 1. DiffusionModelBase.ReflectInstanceDisposables: dict-held disposables were silently leaked. Dictionary<TKey, IDisposable> yields KeyValuePair<,> structs (not IDisposable) when enumerated, so the generic IEnumerable branch never reached the payload. Added an IDictionary branch BEFORE the IEnumerable branch that walks DictionaryEntry and extracts .Value, so dictionary-keyed disposables cascade through Dispose. 2. NoisePredictorBase.ReflectInstanceLayers: identical bug for layer enumeration. Dictionary<TKey, ILayer<T>> was silently dropping layers because the enumerator yields KeyValuePair<,> structs. Same fix — IDictionary branch before IEnumerable branch, extracting Value and matching it against ILayer<T>. 3. CompiledModelHost.Predict (structureVersion invalidation): previous code invalidated the cache IN PLACE when structureVersion changed, then reused the same CompiledModelCache<T> instance. A racing Predict that had just released the lock with an older version could repopulate the same cache after the switch, leaking version-N plans into version-N+1 callers. Now we replace the cache reference with a fresh instance per structure version — every per-version local `cache` closes over a unique object, so once we leave the lock it's frozen to that version. 4. CompiledModelHost.Predict (narrow catch): the fallback-to-eager path previously used catch-all `catch (Exception)`, which would swallow fatal CLR exceptions (OutOfMemoryException, StackOverflowException, AccessViolationException, thread aborts) and then immediately invoke eagerForward() on a poisoned process. Replaced with an exception filter that propagates fatal exceptions and only recovers from compile/replay failures. Follows the same exclusion list the runtime team uses for "recoverable vs. non-recoverable" distinctions. 5. CompiledModelHostTests: 9 new unit tests covering the host's invalidation contract, Dispose idempotency, eager fallback paths (compilation disabled, post-Dispose), structureVersion-bump cache invalidation, explicit Invalidate(), null-guard on eagerForward, concurrent Predict+Invalidate, and concurrent Predict+Dispose. All tests passing locally on net10.0. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * fix: address PR #1143 round-5 — in-flight coordination + empty EnumerateLayers default Two reviewer concerns flagged on the round-4 changes: 1. CompiledModelHost's Dispose/Invalidate + version-swap path could tear down a cache that an in-flight Predict call was still using outside the lock. Add proper in-flight coordination: - _activeCalls counter (incremented inside lock right before release, decremented in Predict's finally block via DrainPendingDisposals). - On version swap, the detached old cache is parked on _pendingDisposeCaches instead of being Disposed inline — the last Predict to finish drains and disposes the queue. - Dispose sets _disposeRequested; if _activeCalls == 0 it immediately cleans up, otherwise defers to DrainPendingDisposals. - Invalidate() follows the same pattern: if in-flight callers hold the cache, park it on the queue instead of disposing inline. - TraceWarning now logs ex.ToString() so stack traces and inner exceptions survive for production diagnostics. 2. NoisePredictorBase.EnumerateLayers() default changes back to Enumerable.Empty<ILayer<T>>() to match the original PR contract. Reflecting by default would dispose layers the predictor doesn't own (injected cross-attention layers from a shared encoder, an injected VAE, etc.). Ownership should be expressed by what the predictor chooses to enumerate, not by what reflection happens to find. Make ReflectInstanceLayers protected so subclasses can opt in explicitly when they DO want the walk. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * fix: address PR #1143 round-5 — in-flight coordination + empty EnumerateLayers default Two reviewer concerns flagged on the round-4 changes: 1. CompiledModelHost's Dispose/Invalidate + version-swap path could tear down a cache that an in-flight Predict call was still using outside the lock. Add proper in-flight coordination: - _activeCalls counter (incremented inside lock right before release, decremented in Predict's finally block via DrainPendingDisposals). - On version swap, the detached old cache is parked on _pendingDisposeCaches instead of being Disposed inline — the last Predict to finish drains and disposes the queue. - Dispose sets _disposeRequested; if _activeCalls == 0 it immediately cleans up, otherwise defers to DrainPendingDisposals. - Invalidate() follows the same pattern: if in-flight callers hold the cache, park it on the queue instead of disposing inline. - TraceWarning now logs ex.ToString() so stack traces and inner exceptions survive for production diagnostics. 2. NoisePredictorBase.EnumerateLayers() default changes back to Enumerable.Empty<ILayer<T>>() to match the original PR contract. Reflecting by default would dispose layers the predictor doesn't own (injected cross-attention layers from a shared encoder, an injected VAE, etc.). Ownership should be expressed by what the predictor chooses to enumerate, not by what reflection happens to find. Make ReflectInstanceLayers protected so subclasses can opt in explicitly when they DO want the walk. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * fix: IDiffusionModel : IDisposable + using var in DiffusionModelTestBase (#1136) ## Summary Completes the Dispose cascade wired up in this PR by exposing `IDisposable` on the `IDiffusionModel<T>` interface itself and flipping the 13 `var model = CreateModel()` calls in `DiffusionModelTestBase` to `using var model`. Without this pair of changes the `DiffusionModelBase<T> : IDisposable` implementation was unreachable through the interface — consumers who held `IDiffusionModel<T>` references had no compile-time visibility of the Dispose contract, so the reflection-walk Dispose that this PR adds was effectively opt-in via an `IDisposable` cast. ## Why this matters for CI (issue #1136) The three cancelled Diffusion ModelFamily shards (A-I / J-R / S-Z) on master baseline [24398739627](https://github.com/ooples/AiDotNet/actions/runs/24398739627) were hitting the 45-minute wall-clock budget with ~950 `OutOfMemoryException`s across ~255 diffusion model tests. Every test was constructing a production-default model (DiT-XL: ~4 GB eagerly-allocated weights) and letting the reference fall out of scope without disposing. With no Dispose path, the rented weight tensors stayed live in the TensorAllocator pool — and the next test's allocation stacked on top. `using var model = CreateModel()` now calls Dispose() at the end of every test method, which walks the reflection-discovered child disposables added in this PR (noise predictors, schedulers) and releases their layer trees back to the pool. Memory now plateaus across tests instead of monotonically climbing to OOM. ## Scope - `IDiffusionModel<T>` inherits `System.IDisposable`. `DiffusionModelBase<T>` already implemented `IDisposable` in this PR; this commit just exposes it on the interface so the compiler can see it. - `DiffusionModelTestBase`: 13 call sites converted to `using var model`. `Clone()` call in `Clone_ShouldProduceIdenticalOutput` left unchanged — it returns `IFullModel<...>` which doesn't derive from `IDisposable`. The primary `model` is now disposed, which is the dominant memory path. ## Not in scope - `NeuralNetworkModelTestBase` — already handled on PR #1137 branch via the `INeuralNetwork<T> : IDisposable` change. - Other ModelFamily test bases (20+) use different model interfaces (IClassificationModel, ITimeSeriesModel, etc.). Adding `IDisposable` to each is the natural follow-up once their respective base classes implement Dispose — tracked separately, not this PR. Build clean on net10.0 + net471, 0 errors. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * fix: IAsyncLifetime forced GC between DiffusionModelTestBase tests (#1136) ## Summary Completes Part 4 of the #1136 fix plan for DiffusionModelTestBase by implementing `IAsyncLifetime.DisposeAsync` with the standard two-pass GC pattern (`GC.Collect → WaitForPendingFinalizers → GC.Collect`) that runs AFTER each test disposes its model via `using var model = CreateModel()`. ## Why a forced GC is necessary The prior commit (<a2acab430>) added `using var model = CreateModel()` to all 13 Diffusion test methods, which calls Dispose at end-of-test — the Dispose walks child layers and returns their rented weight buffers to the TensorAllocator pool. But `Dispose()` only releases the REFERENCES; the managed weight arrays (300 MB per DiT-XL model instance) sit in gen-2 heap until GC runs a compacting pass. On 16 GB Windows CI runners, GC waits for memory pressure — but the ~255 sequential diffusion tests allocate fast enough to hit the 45-min wall clock before gen-2 pressure triggers a collection. The forced GC runs explicitly after each test disposes, reclaiming the pool-returned arrays immediately. Memory now plateaus across tests instead of monotonically climbing to OOM. ## Pattern Implements `IAsyncLifetime` (xunit's per-test lifecycle hook): - `InitializeAsync`: no-op — no ambient state to set up - `DisposeAsync`: two-pass GC (the standard .NET pattern for forcing finalizers to run AND reclaiming the collected memory) ## Scope DiffusionModelTestBase only. The same pattern should apply to every ModelFamily shard (NeuralNetworkModelTestBase, AudioDiffusionTestBase, VideoDiffusionTestBase, etc.), but those belong in separate commits as each family matures its own Dispose cascade. Build clean on net10.0, 0 errors. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * perf: compact LOH in DiffusionModelTestBase teardown (#1136) Upgrades the IAsyncLifetime DisposeAsync hook added in the prior commit to use a blocking COMPACTING Gen-2 collection with explicit Large Object Heap compaction, matching the pattern applied on PR #1148's DiffusionUnitTestBase. Diffusion model weight tensors are typically several hundred MB each, well above the 85KB LOH threshold. Plain `GC.Collect` sweeps the LOH but does NOT compact it — across ~255 sequential diffusion tests, LOH fragmentation accumulates until the next allocation can't find a contiguous region even though total free bytes remain large, producing `OutOfMemoryException` on a runner that APPEARS to have plenty of memory. Setting `GCLargeObjectHeapCompactionMode.CompactOnce` on the next Gen-2 pass forces LOH compaction; the mode auto-resets to Default after each use, so this is scoped per-teardown — no process-wide side-effect. Same fix applied on PR #1148 (b4aa53f) per CodeRabbit review feedback. Applying here too for consistency — both test bases run diffusion models with similar LOH fragmentation risk. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * perf: serialize LOH compaction with static lock (#1136) Matches the lock added on PR #1148 (6794492). xunit parallelizes across test-classes by default — two derived test classes can hit DisposeAsync concurrently on different threads, racing on the process-global GCSettings.LargeObjectHeapCompactionMode flag. Wrapping the whole GC sequence in a static lock keeps LOH compaction deterministic per teardown. Perf note: the lock only serializes teardowns, not the tests themselves. The GC.Collect calls are already blocking anyway, so the extra lock contention is negligible compared to the collection cost. Build clean on net10.0. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Co-authored-by: franklinic <franklin@ivorycloud.com>
IDiffusionModel.cs: both sides add IDisposable — take master's fully-qualified System.IDisposable on its own line (cleaner formatting). DiffusionModelBase.cs: PR #1140 had a simple Dispose(bool) stub; master (via merged PR #1143) has the full reflection-based EnumerateDisposableComponents + IDictionary walk + Dispose cascade. Take master's complete version — it supersedes the stub. DiffusionModelTestBase.cs: combine both sides — - Take master's IAsyncLifetime <remarks> doc block (PR #1140 branch predates the IAsyncLifetime addition). - Keep PR #1140's ParameterCount > 0 optimization (lazy-friendly, avoids forcing 4 GB weight materialization just for an existence check) over master's GetParameters().Length > 0 (the old eager version). Build clean on net10.0. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 7
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
src/Diffusion/NoisePredictors/DiTNoisePredictor.cs (1)
451-455:⚠️ Potential issue | 🔴 CriticalInitialize lazy layers in
PredictNoiseWithEmbedding().This entry point now bypasses
EnsureLayersInitialized(), so calling it on a freshDiTNoisePredictor<T>throws"Layers not initialized."before the first forward pass.Proposed fix
public override Tensor<T> PredictNoiseWithEmbedding(Tensor<T> noisySample, Tensor<T> timeEmbedding, Tensor<T>? conditioning = null) { + EnsureLayersInitialized(); _lastInput = noisySample; return Forward(noisySample, timeEmbedding, conditioning); }🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/Diffusion/NoisePredictors/DiTNoisePredictor.cs` around lines 451 - 455, PredictNoiseWithEmbedding currently sets _lastInput and calls Forward directly, which can throw "Layers not initialized."; update PredictNoiseWithEmbedding in DiTNoisePredictor<T> to call EnsureLayersInitialized() at the start (before accessing _lastInput or calling Forward) so lazy layers are initialized, then proceed to set _lastInput and call Forward(noisySample, timeEmbedding, conditioning).
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@src/NeuralNetworks/Layers/MultiHeadAttentionLayer.cs`:
- Around line 344-347: Validate headCount and embeddingDimension divisibility
before computing _headDimension in the MultiHeadAttentionLayer constructors:
check that headCount > 0 and that embeddingDimension % headCount == 0, and throw
an ArgumentOutOfRangeException/ArgumentException with a clear message if either
check fails; apply the same validation in both constructor locations where
_headCount, _headDimension and _embeddingDimension are set (the assignment
blocks around _headCount = headCount / _headDimension = embeddingDimension /
headCount / _embeddingDimension = embeddingDimension).
- Around line 349-379: The eager branch currently always calls
InitializeParameters(), which applies the default Xavier init instead of
honoring a provided non-lazy IInitializationStrategy<T>; change the eager path
so that when initializationStrategy is non-null and
initializationStrategy.IsLazy == false you invoke the strategy's initialization
(e.g., call the strategy's Initialize/Apply method on _queryWeights,
_keyWeights, _valueWeights, _outputWeights, _outputBias) instead of
InitializeParameters(), then RegisterTrainableParameter(...) as before and set
_isInitialized=true; replicate the exact same change for the other eager block
referenced (the one around lines 408-436) so both eager initializations use the
supplied IInitializationStrategy<T>.
- Around line 895-901: ForwardGpu currently runs projection kernels on
unmaterialized Q/K/V/O tensors for lazy-initialized MultiHeadAttentionLayer
instances; update the start of ForwardGpu in class MultiHeadAttentionLayer to
call EnsureWeightsAllocated() and EnsureInitialized() (same as
ForwardInternal()) before any projection kernel invocation so the weight tensors
are materialized and sub-layers (ROPE/ALiBi) are registered.
In `@src/NeuralNetworks/Layers/SelfAttentionLayer.cs`:
- Around line 394-420: The eager constructor branch currently calls
InitializeParameters() and registers weights unconditionally, ignoring the
provided initializationStrategy; change the logic so you only call
InitializeParameters() and RegisterTrainableParameter(...) and set
_isInitialized=true when initializationStrategy indicates non-lazy
initialization (i.e., initializationStrategy is not { IsLazy: true }), otherwise
leave tensors as placeholders and _isInitialized=false; apply the same fix to
the other eager branch (the similar block around the other constructor/branch
referenced at lines ~479-501) so the initializationStrategy parameter is honored
in both places.
- Around line 539-542: The GPU forward path in SelfAttentionLayer (ForwardGpu)
doesn't materialize lazy-initialized weights, so it may use zero-sized
placeholders; call EnsureInitialized() at the start of ForwardGpu (same as CPU
forward) to force creation of Q/K/V and bias tensors before any GPU
projections/BLAS calls, and then proceed to use the now-populated member tensors
(Q/K/V/bias) so the first GPU inference/training no longer fails.
In `@src/NeuralNetworks/NeuralNetworkBase.cs`:
- Around line 4634-4654: The second cascade that iterates Layers and calls
Dispose on IDisposable instances is duplicating disposals already handled by
DisposeOnceGuard.TryDispose(...) and can double-return pooled buffers (e.g., via
TensorAllocator), so remove this extra foreach block that checks Layers and
calls disposable.Dispose(); ensure only DisposeOnceGuard.TryDispose is used for
layer disposal to prevent double-dispose of tensors and allocator corruption
(refer to Layers, DisposeOnceGuard.TryDispose, IDisposable.Dispose in
NeuralNetworkBase).
In `@tests/AiDotNet.Tests/ModelFamilyTests/Base/DiffusionModelTestBase.cs`:
- Around line 87-104: The class DiffusionModelTestBase contains duplicate
IAsyncLifetime implementations: the methods InitializeAsync() and DisposeAsync()
shown here are already declared earlier in the class; remove this second pair
(or consolidate with the earlier implementations) so there are no duplicate
member definitions—specifically delete or merge the redundant InitializeAsync()
and DisposeAsync() methods in DiffusionModelTestBase.
---
Outside diff comments:
In `@src/Diffusion/NoisePredictors/DiTNoisePredictor.cs`:
- Around line 451-455: PredictNoiseWithEmbedding currently sets _lastInput and
calls Forward directly, which can throw "Layers not initialized."; update
PredictNoiseWithEmbedding in DiTNoisePredictor<T> to call
EnsureLayersInitialized() at the start (before accessing _lastInput or calling
Forward) so lazy layers are initialized, then proceed to set _lastInput and call
Forward(noisySample, timeEmbedding, conditioning).
🪄 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: 834c6005-e2d1-441c-a0ce-c21449c0baf4
📒 Files selected for processing (20)
src/Diffusion/NoisePredictors/AsymmDiTPredictor.cssrc/Diffusion/NoisePredictors/DiTNoisePredictor.cssrc/Diffusion/NoisePredictors/DiffusionResBlock.cssrc/Diffusion/NoisePredictors/EMMDiTPredictor.cssrc/Diffusion/NoisePredictors/FlagDiTPredictor.cssrc/Diffusion/NoisePredictors/FluxDoubleStreamPredictor.cssrc/Diffusion/NoisePredictors/MMDiTNoisePredictor.cssrc/Diffusion/NoisePredictors/MMDiTXNoisePredictor.cssrc/Diffusion/NoisePredictors/NoisePredictorBase.cssrc/Diffusion/NoisePredictors/SiTPredictor.cssrc/Diffusion/NoisePredictors/UNetNoisePredictor.cssrc/Diffusion/NoisePredictors/UViTNoisePredictor.cssrc/Diffusion/NoisePredictors/VideoUNetPredictor.cssrc/Interfaces/INeuralNetworkModel.cssrc/NeuralNetworks/Layers/ConvolutionalLayer.cssrc/NeuralNetworks/Layers/MultiHeadAttentionLayer.cssrc/NeuralNetworks/Layers/SelfAttentionLayer.cssrc/NeuralNetworks/NeuralNetworkBase.cstests/AiDotNet.Tests/ModelFamilyTests/Base/DiffusionModelTestBase.cstests/AiDotNet.Tests/ModelFamilyTests/Base/NeuralNetworkModelTestBase.cs
ooples
added a commit
that referenced
this pull request
Apr 16, 2026
* perf: lazy weight init in diffusion noise predictors (part 1/5 of #1136) Every diffusion model test on CI was OOMing during construction because DiT/MMDiT/UNet noise predictors eagerly allocate ~4 GB of weight tensors in their ctors — `new DenseLayer<T>(...)` calls `TensorAllocator.Rent<T>` before the model has even seen an input. With 255 diffusion models × ~10 tests each running in a shared xunit process, sequential tests stack up and OOM on the 16 GB Windows CI runners, cancelling all three Diffusion ModelFamily shards at the 45-minute wall clock. This part of the fix makes model construction O(1) by threading `InitializationStrategies<T>.Lazy` into every internal Dense/Conv layer built by a noise predictor. Weight tensors stay at shape [0,0] until the first Forward() pass actually needs them. Adds two protected helpers on NoisePredictorBase<T>: - LazyDense(int, int, IActivationFunction?) — lazy dense layer - LazyDenseVec(int, int, IVectorActivationFunction) — lazy dense for vector activations (distinct name avoids ctor-overload ambiguity on activations that implement both interfaces) - LazyConv2D(...) — lazy 2D convolutional layer Converts all 11 noise predictors + DiffusionResBlock: - DiTNoisePredictor (patch/time embeds + final + MLP/attention blocks) - MMDiTNoisePredictor (image/text streams + joint + single blocks) - MMDiTXNoisePredictor - EMMDiTPredictor - FlagDiTPredictor - FluxDoubleStreamPredictor (double + single streams) - AsymmDiTPredictor - SiTPredictor - UViTNoisePredictor (encoder + decoder + skip projections) - UNetNoisePredictor (downsample) - VideoUNetPredictor (time embed + spatial/temporal res blocks + in/out conv) - DiffusionResBlock (norm/conv1 + timeMlp + norm/conv2 + skip conv) Parts 2-5 (MHA/LayerNorm lazy init, Dispose→pool return, test lifecycle, GetParameters→ParameterCount swap) follow in subsequent commits. Refs #1136 * perf: lazy weight init in diffusion noise predictors (part 1/5 of #1136) Every diffusion model test on CI was OOMing during construction because DiT/MMDiT/UNet noise predictors eagerly allocate ~4 GB of weight tensors in their ctors — `new DenseLayer<T>(...)` calls `TensorAllocator.Rent<T>` before the model has even seen an input. With 255 diffusion models × ~10 tests each running in a shared xunit process, sequential tests stack up and OOM on the 16 GB Windows CI runners, cancelling all three Diffusion ModelFamily shards at the 45-minute wall clock. This part of the fix makes model construction O(1) by threading `InitializationStrategies<T>.Lazy` into every internal Dense/Conv layer built by a noise predictor. Weight tensors stay at shape [0,0] until the first Forward() pass actually needs them. Adds two protected helpers on NoisePredictorBase<T>: - LazyDense(int, int, IActivationFunction?) — lazy dense layer - LazyDenseVec(int, int, IVectorActivationFunction) — lazy dense for vector activations (distinct name avoids ctor-overload ambiguity on activations that implement both interfaces) - LazyConv2D(...) — lazy 2D convolutional layer Converts all 11 noise predictors + DiffusionResBlock: - DiTNoisePredictor (patch/time embeds + final + MLP/attention blocks) - MMDiTNoisePredictor (image/text streams + joint + single blocks) - MMDiTXNoisePredictor - EMMDiTPredictor - FlagDiTPredictor - FluxDoubleStreamPredictor (double + single streams) - AsymmDiTPredictor - SiTPredictor - UViTNoisePredictor (encoder + decoder + skip projections) - UNetNoisePredictor (downsample) - VideoUNetPredictor (time embed + spatial/temporal res blocks + in/out conv) - DiffusionResBlock (norm/conv1 + timeMlp + norm/conv2 + skip conv) Parts 2-5 (MHA/LayerNorm lazy init, Dispose→pool return, test lifecycle, GetParameters→ParameterCount swap) follow in subsequent commits. Refs #1136 * perf: lazy-friendly Parameters_ShouldBeNonEmpty + GC between tests (parts 4+5/5 of #1136) Two follow-on pieces to the lazy-init-in-noise-predictors change (PR 1 of the #1136 series). Together with that first commit they should clear the three cancelled Diffusion ModelFamily shards on CI — if not fully, they put the next diagnostics in reach instead of hiding behind OOM. Part 5: swap `GetParameters().Length > 0` for `ParameterCount > 0`. The `Parameters_ShouldBeNonEmpty` tests in DiffusionModelTestBase and NeuralNetworkModelTestBase are just asking "does this model have any learnable parameters?" — a question `ParameterCount` answers in O(1) without materializing the full flattened parameter vector. Calling GetParameters() on a lazily-constructed DiT-XL forces every DenseLayer to eagerly allocate weights (EnsureInitialized cascades down) just to count them — that alone reproduces the ~4 GB OOM even with lazy init in place. ParameterCount propagates through the layer list without triggering lazy materialization, so the existence check stays cheap. Semantically identical assertion, not a weakening — both check that the model has learnable parameters. Just using the right API. Part 4: IAsyncLifetime.DisposeAsync forces a full GC cycle between test methods in both ModelFamily bases. Each test instantiates a fresh production-sized model (VGG16BN, DiT-XL, SDXL, etc.). Without explicit GC pressure, the shared xunit process holds onto previous test's weight tensors until the collector runs on its own schedule. Sequential tests in a shard stack up live weight allocations that exceed the 16 GB Windows runner budget. IAsyncLifetime.DisposeAsync runs after every [Fact] in a class, so a Collect/WaitForPendingFinalizers/Collect cycle there releases the previous test's weight tensors before the next test's model constructs. Combined with Part 1 (lazy construction) and Part 5 (no forced materialization on existence check), the working-set for a single test stays bounded to what that test actually touches. Parts 2 (MHA/LayerNorm lazy init) and 3 (Dispose → pool-return) are deferred to a follow-up PR once we see how much of the OOM is addressed by Parts 1/4/5 alone. Refs #1136 * perf: lazy-friendly Parameters_ShouldBeNonEmpty + GC between tests (parts 4+5/5 of #1136) Two follow-on pieces to the lazy-init-in-noise-predictors change (PR 1 of the #1136 series). Together with that first commit they should clear the three cancelled Diffusion ModelFamily shards on CI — if not fully, they put the next diagnostics in reach instead of hiding behind OOM. Part 5: swap `GetParameters().Length > 0` for `ParameterCount > 0`. The `Parameters_ShouldBeNonEmpty` tests in DiffusionModelTestBase and NeuralNetworkModelTestBase are just asking "does this model have any learnable parameters?" — a question `ParameterCount` answers in O(1) without materializing the full flattened parameter vector. Calling GetParameters() on a lazily-constructed DiT-XL forces every DenseLayer to eagerly allocate weights (EnsureInitialized cascades down) just to count them — that alone reproduces the ~4 GB OOM even with lazy init in place. ParameterCount propagates through the layer list without triggering lazy materialization, so the existence check stays cheap. Semantically identical assertion, not a weakening — both check that the model has learnable parameters. Just using the right API. Part 4: IAsyncLifetime.DisposeAsync forces a full GC cycle between test methods in both ModelFamily bases. Each test instantiates a fresh production-sized model (VGG16BN, DiT-XL, SDXL, etc.). Without explicit GC pressure, the shared xunit process holds onto previous test's weight tensors until the collector runs on its own schedule. Sequential tests in a shard stack up live weight allocations that exceed the 16 GB Windows runner budget. IAsyncLifetime.DisposeAsync runs after every [Fact] in a class, so a Collect/WaitForPendingFinalizers/Collect cycle there releases the previous test's weight tensors before the next test's model constructs. Combined with Part 1 (lazy construction) and Part 5 (no forced materialization on existence check), the working-set for a single test stays bounded to what that test actually touches. Parts 2 (MHA/LayerNorm lazy init) and 3 (Dispose → pool-return) are deferred to a follow-up PR once we see how much of the OOM is addressed by Parts 1/4/5 alone. Refs #1136 * fix: address pr #1137 review comments — film conditioning, temporal mixing, lazy guards VideoUNet timestep conditioning (Ho et al. 2022 "Video Diffusion Models" §3.1): * Each VideoBlock now has a TimeCondProjection (DenseLayer: timeEmbedDim → C*2) * New ApplyFiLMConditioning: x = x * (1 + scale) + shift, broadcast over spatial dims * Previously timeEmbed was dropped on the floor — output was invariant to timestep VideoUNet temporal residual processing (Ho et al. 2022 §3.1): * ApplyTemporalProcessing reshapes [B,C,F,H,W] → [B*C*H*W, F], applies temporal DenseLayer, reshapes back, adds residual. Was a no-op clone. * CreateTemporalResBlock now uses DenseLayer(numFrames, numFrames) for temporal mixing DiT: Forward() calls EnsureLayersInitialized() for lazy-init guard. Tests: using var for IDisposable networks, comment for non-disposable diffusion models. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * fix: PR #1137 review comments — FiLM conditioning, temporal mixing, lazy guards VideoUNet timestep conditioning (Ho et al. 2022 "Video Diffusion Models" §3.1): * Each VideoBlock now has a TimeCondProjection (DenseLayer: timeEmbedDim → C*2) * New ApplyFiLMConditioning: x = x * (1 + scale) + shift, broadcast over spatial dims * Previously timeEmbed was dropped on the floor — output was invariant to timestep VideoUNet temporal residual processing (Ho et al. 2022 §3.1): * ApplyTemporalProcessing reshapes [B,C,F,H,W] → [B*C*H*W, F], applies temporal DenseLayer, reshapes back, adds residual. Was a no-op clone. * CreateTemporalResBlock now uses DenseLayer(numFrames, numFrames) for temporal mixing DiT: Forward() calls EnsureLayersInitialized() for lazy-init guard. Tests: using var for IDisposable networks, comment for non-disposable diffusion models. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * fix: address pr #1137 round-2 — thread-safe dit init + remove duplicate xml summaries DiT: EnsureLayersInitialized now uses double-checked locking with a volatile flag so concurrent first calls can't duplicate _blocks entries. Matches the pattern used by FeedForwardLayer/ConvolutionalLayer. VideoUNet: removed stale duplicate <summary> blocks above ApplyVideoBlock and ApplyTemporalProcessing that were left behind when the expanded documentation was added in the prior commit. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * fix: PR #1137 round-2 — thread-safe DiT init + remove duplicate XML summaries DiT: EnsureLayersInitialized now uses double-checked locking with a volatile flag so concurrent first calls can't duplicate _blocks entries. Matches the pattern used by FeedForwardLayer/ConvolutionalLayer. VideoUNet: removed stale duplicate <summary> blocks above ApplyVideoBlock and ApplyTemporalProcessing that were left behind when the expanded documentation was added in the prior commit. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * fix: address PR #1137 review feedback — VideoUNet correctness + lazy-init retry-safety Critical: - ApplyFiLMConditioning splits condVec via Engine.TensorSlice + Engine.Reshape instead of AsSpan + new Tensor<T>. The previous implementation copied data into raw managed arrays then rebuilt fresh tensors, severing the autograd tape — gradients couldn't flow back into the TimeCondProjection weights and the FiLM layer never learned. - TimeCondProjection now wired into AddBlockParameters / SetBlockParameters / AddBlockGradients with consistent ordering (between CrossAttention and Downsample). Previously these new trainable layers were ignored by the param/grad/serialize plumbing — weights silently missed optimizer updates, didn't survive Save/Load/Clone round-trips, and gradient maps misaligned. - ApplyTemporalProcessing permutes [B,C,F,H,W] → [B,C,H,W,F] BEFORE flattening to [B*C*H*W, F], then permutes back. The previous direct Reshape grouped contiguous spatial elements into the F slot instead of one temporal vector per (b,c,h,w) position — the temporal Dense was mixing spatial neighbors instead of frames. - ApplyTemporalProcessing also clones video._shape before reuse — reusing the tensor's mutable backing array could let downstream mutations corrupt the source tensor's shape. Major: - DiTNoisePredictor.EnsureLayersInitialized clears _blocks before InitializeLayers so a partial-init failure followed by a retry doesn't append a second set of blocks on top of the partial state. _layersInitialized remains the LAST step inside the lock so observers never see a half-built graph. - Parameters_ShouldBeNonEmpty in DiffusionModelTestBase + NeuralNetworkModelTestBase uses ParameterCount instead of GetParameters().Length. The previous test forced full lazy-init allocation just to measure the length, defeating the PR's lazy-allocation work — ParameterCount asks the same semantic question via the lazy-friendly metadata API. Build clean on net10.0 + net471, 0 errors. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * fix: address PR #1137 review feedback — VideoUNet correctness + lazy-init retry-safety Critical: - ApplyFiLMConditioning splits condVec via Engine.TensorSlice + Engine.Reshape instead of AsSpan + new Tensor<T>. The previous implementation copied data into raw managed arrays then rebuilt fresh tensors, severing the autograd tape — gradients couldn't flow back into the TimeCondProjection weights and the FiLM layer never learned. - TimeCondProjection now wired into AddBlockParameters / SetBlockParameters / AddBlockGradients with consistent ordering (between CrossAttention and Downsample). Previously these new trainable layers were ignored by the param/grad/serialize plumbing — weights silently missed optimizer updates, didn't survive Save/Load/Clone round-trips, and gradient maps misaligned. - ApplyTemporalProcessing permutes [B,C,F,H,W] → [B,C,H,W,F] BEFORE flattening to [B*C*H*W, F], then permutes back. The previous direct Reshape grouped contiguous spatial elements into the F slot instead of one temporal vector per (b,c,h,w) position — the temporal Dense was mixing spatial neighbors instead of frames. - ApplyTemporalProcessing also clones video._shape before reuse — reusing the tensor's mutable backing array could let downstream mutations corrupt the source tensor's shape. Major: - DiTNoisePredictor.EnsureLayersInitialized clears _blocks before InitializeLayers so a partial-init failure followed by a retry doesn't append a second set of blocks on top of the partial state. _layersInitialized remains the LAST step inside the lock so observers never see a half-built graph. - Parameters_ShouldBeNonEmpty in DiffusionModelTestBase + NeuralNetworkModelTestBase uses ParameterCount instead of GetParameters().Length. The previous test forced full lazy-init allocation just to measure the length, defeating the PR's lazy-allocation work — ParameterCount asks the same semantic question via the lazy-friendly metadata API. Build clean on net10.0 + net471, 0 errors. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * fix: address PR #1137 round-3 — defensive list copy + LazyDense for video factories - DiTNoisePredictor constructor snapshots customBlocks via defensive copy. With lazy init, the caller-owned List<DiTBlock> is read at first-use instead of construction time, so a caller mutating it after 'new DiTNoisePredictor(...)' would silently change the model's block graph. Snapshotting at construction preserves construction-time semantics. - VideoUNetPredictor.CreateTemporalResBlock and CreateTimeCondProjection switch to LazyDense. These factories were called once per encoder + middle + decoder block — eager DenseLayer allocation re-introduced the OOM pressure the PR's lazy-init was specifically designed to relieve. Build clean on net10.0 + net471, 0 errors. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * fix: address PR #1137 round-3 — defensive list copy + LazyDense for video factories - DiTNoisePredictor constructor snapshots customBlocks via defensive copy. With lazy init, the caller-owned List<DiTBlock> is read at first-use instead of construction time, so a caller mutating it after 'new DiTNoisePredictor(...)' would silently change the model's block graph. Snapshotting at construction preserves construction-time semantics. - VideoUNetPredictor.CreateTemporalResBlock and CreateTimeCondProjection switch to LazyDense. These factories were called once per encoder + middle + decoder block — eager DenseLayer allocation re-introduced the OOM pressure the PR's lazy-init was specifically designed to relieve. Build clean on net10.0 + net471, 0 errors. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * fix: address PR #1137 round-4 — numClasses no-op guard, VideoUNet lazy-allocs, internal shape access Critical: - DiTNoisePredictor constructor now throws NotSupportedException when numClasses > 0. The _labelEmbed layer is created and participates in param/grad/serialize plumbing but no forward path actually injects the class index — weights would train against zero-gradient signal silently. Every caller in the codebase currently passes numClasses=0 (the feature has zero active users), so the throw is safe. The plumbing stays so a future PR can wire the class-injection path end-to-end (following DiT paper §3.2: fuse label embedding into projected time embedding before AdaLN modulation). Major: - VideoUNetPredictor._inputConv, _outputConv, _timeEmbedMlp1, _timeEmbedMlp2, _imageCondProjection, CreateSpatialResBlock, and CreateDownsample now use the Lazy factories. Previously only 2 of ~10 layer factories were lazy, re-introducing the exact OOM pressure the PR was meant to relieve. MHA attention layers and DeconvolutionalLayer don't have Lazy factories on this branch yet (attention lazy support is on PR #1140; no LazyDeconv exists), so those remain eager — documented as a follow-up. - ApplyFiLMConditioning validates projection output channels match x's channel dim. Without this check, a misconfigured projection silently produces wrong-shape broadcasts or broadcasts into the wrong axis. - ApplyTemporalProcessing and ApplyTemporalAttention now use video.Shape (public) instead of video._shape (internal backing field). Decouples the predictor from Tensor<T>'s storage layout. Build clean on net10.0 + net471, 0 errors. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * fix: address PR #1137 round-4 — numClasses no-op guard, VideoUNet lazy-allocs, internal shape access Critical: - DiTNoisePredictor constructor now throws NotSupportedException when numClasses > 0. The _labelEmbed layer is created and participates in param/grad/serialize plumbing but no forward path actually injects the class index — weights would train against zero-gradient signal silently. Every caller in the codebase currently passes numClasses=0 (the feature has zero active users), so the throw is safe. The plumbing stays so a future PR can wire the class-injection path end-to-end (following DiT paper §3.2: fuse label embedding into projected time embedding before AdaLN modulation). Major: - VideoUNetPredictor._inputConv, _outputConv, _timeEmbedMlp1, _timeEmbedMlp2, _imageCondProjection, CreateSpatialResBlock, and CreateDownsample now use the Lazy factories. Previously only 2 of ~10 layer factories were lazy, re-introducing the exact OOM pressure the PR was meant to relieve. MHA attention layers and DeconvolutionalLayer don't have Lazy factories on this branch yet (attention lazy support is on PR #1140; no LazyDeconv exists), so those remain eager — documented as a follow-up. - ApplyFiLMConditioning validates projection output channels match x's channel dim. Without this check, a misconfigured projection silently produces wrong-shape broadcasts or broadcasts into the wrong axis. - ApplyTemporalProcessing and ApplyTemporalAttention now use video.Shape (public) instead of video._shape (internal backing field). Decouples the predictor from Tensor<T>'s storage layout. Build clean on net10.0 + net471, 0 errors. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * fix: address pr #1137 round-3 — full dit class conditioning, videounet film/lazy Addresses 4 unresolved PR #1137 review comments: 1. VideoUNet FiLM validation — use x.Shape as ground truth for batch/channels; validate condVec.Shape[0] == batchSize and condVec.Shape[^1] == channels*2 with descriptive ArgumentException on mismatch. 2. Public Shape API — ApplyTemporalProcessing now uses video.Shape.ToArray() instead of internal video._shape backing field. 3. Full DiT class conditioning per Peebles & Xie 2022 §3.2 — _labelEmbed output changed from _hiddenSize to timeEmbedDim so it can be summed with timeEmbed. When numClasses > 0 AND conditioning is supplied, Forward() interprets conditioning as one-hot class labels [B, numClasses], projects them to timeEmbedDim, adds to timeEmbed, and feeds the combined adaLnEmbed to all AdaLN modulation paths (block-level and FinalLayerWithAdaLN). Validates conditioning.Shape[^1] == numClasses. Class-conditioned mode does NOT feed conditioning to cross-attention (exclusive paths). 4. Full lazy VideoUNet conversion — converted all DenseLayer/ConvolutionalLayer factories to LazyDense/LazyConv2D helpers: _inputConv, _outputConv, _timeEmbedMlp1/2, _imageCondProjection, CreateSpatialResBlock, CreateTemporalResBlock (now LazyDense([F,F]) for temporal axis mixing), CreateTimeCondProjection, CreateDownsample. MultiHeadAttentionLayer (spatial/temporal/cross attention) and DeconvolutionalLayer (CreateUpsample) remain eager — no lazy variants exist for those layer types yet. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * fix: PR #1137 round-3 — full DiT class conditioning, VideoUNet FiLM/lazy Addresses 4 unresolved PR #1137 review comments: 1. VideoUNet FiLM validation — use x.Shape as ground truth for batch/channels; validate condVec.Shape[0] == batchSize and condVec.Shape[^1] == channels*2 with descriptive ArgumentException on mismatch. 2. Public Shape API — ApplyTemporalProcessing now uses video.Shape.ToArray() instead of internal video._shape backing field. 3. Full DiT class conditioning per Peebles & Xie 2022 §3.2 — _labelEmbed output changed from _hiddenSize to timeEmbedDim so it can be summed with timeEmbed. When numClasses > 0 AND conditioning is supplied, Forward() interprets conditioning as one-hot class labels [B, numClasses], projects them to timeEmbedDim, adds to timeEmbed, and feeds the combined adaLnEmbed to all AdaLN modulation paths (block-level and FinalLayerWithAdaLN). Validates conditioning.Shape[^1] == numClasses. Class-conditioned mode does NOT feed conditioning to cross-attention (exclusive paths). 4. Full lazy VideoUNet conversion — converted all DenseLayer/ConvolutionalLayer factories to LazyDense/LazyConv2D helpers: _inputConv, _outputConv, _timeEmbedMlp1/2, _imageCondProjection, CreateSpatialResBlock, CreateTemporalResBlock (now LazyDense([F,F]) for temporal axis mixing), CreateTimeCondProjection, CreateDownsample. MultiHeadAttentionLayer (spatial/temporal/cross attention) and DeconvolutionalLayer (CreateUpsample) remain eager — no lazy variants exist for those layer types yet. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * fix: address PR #1137 round-5 — correct temporal mixing axis + clean API Two reviewer concerns on VideoUNetPredictor: 1. ApplyTemporalProcessing was doing a plain reshape from [B,C,F,H,W] to [B*C*H*W, F] and treating each row as a per-(b,c,h,w) frame vector. That's incorrect: F is not the innermost dimension in the source layout, so the reshape produces rows that mix values across H and W as well. The DenseLayer would then learn a cross-spatial mixing rather than temporal mixing, silently corrupting the frame-axis refinement the block is supposed to perform. Fix: permute [B,C,F,H,W] → [B,C,H,W,F] via Engine.TensorPermute first, THEN reshape to [B*C*H*W, F] so rows ARE the (b,c,h,w)-indexed frame vectors. Apply the mixing layer, reshape back to [B,C,H,W,F], then un-permute to [B,C,F,H,W] for the caller. Matches the pattern ApplyTemporalAttention already uses. 2. CreateTemporalResBlock(int channels) ignored its channels parameter — the block is sized by _numFrames, not by the channel count. The unused parameter mislead every caller into sizing the block by channel count. Rename to CreateTemporalMixingBlock() and drop the parameter. Update the four call sites (encoder block, middle block x2, decoder block) to match. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * fix: address PR #1137 round-5 — correct temporal mixing axis + clean API Two reviewer concerns on VideoUNetPredictor: 1. ApplyTemporalProcessing was doing a plain reshape from [B,C,F,H,W] to [B*C*H*W, F] and treating each row as a per-(b,c,h,w) frame vector. That's incorrect: F is not the innermost dimension in the source layout, so the reshape produces rows that mix values across H and W as well. The DenseLayer would then learn a cross-spatial mixing rather than temporal mixing, silently corrupting the frame-axis refinement the block is supposed to perform. Fix: permute [B,C,F,H,W] → [B,C,H,W,F] via Engine.TensorPermute first, THEN reshape to [B*C*H*W, F] so rows ARE the (b,c,h,w)-indexed frame vectors. Apply the mixing layer, reshape back to [B,C,H,W,F], then un-permute to [B,C,F,H,W] for the caller. Matches the pattern ApplyTemporalAttention already uses. 2. CreateTemporalResBlock(int channels) ignored its channels parameter — the block is sized by _numFrames, not by the channel count. The unused parameter mislead every caller into sizing the block by channel count. Rename to CreateTemporalMixingBlock() and drop the parameter. Update the four call sites (encoder block, middle block x2, decoder block) to match. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * fix: address pr #1137 round-5 — axis-correct temporal reshape + tests + film 1d support Addresses 3 unresolved PR #1137 review comments: 1. ApplyTemporalProcessing: axis-correct temporal mixing. The reviewer correctly identified that a plain Reshape from [B, C, F, H, W] to [B*C*H*W, F] does NOT produce per-(b,c,h,w) temporal vectors — F is NOT the innermost dim in that layout, so the reshape enumerator mixes values across H/W into the "F slot". Fixed by permuting [B, C, F, H, W] → [B, C, H, W, F] BEFORE reshape to make F the innermost (row-major last) position, then permute back after the mixing. Mirrors the pattern ApplyTemporalAttention already uses (permute-before-reshape for the same axis-correctness reason). 2. CreateTemporalResBlock: removed misleading unused `channels` parameter. The layer is sized by `_numFrames`, not by channel count. Renamed to CreateTemporalMixingLayer() (per reviewer's suggested rename), removed the unused parameter, and updated all 4 call sites. The new name self-documents that it creates a frame-mixing layer. 3. Integration tests for timestep FiLM + temporal path. Added VideoUNetPredictorIntegrationTests exercising via reflection: - ApplyTemporalProcessing_PreservesShape: end-to-end [B,C,F,H,W] input produces matching-shape output (catches regression in the axis fix). - ApplyTemporalProcessing_MixesAcrossFramesAxis: per-frame-constant video stays per-frame-constant after mixing (catches the axis bug directly — a broken layout would mix H/W into the frame axis, producing non-constant-within-frame outputs). - ApplyFiLMConditioning_ModulatesFeatureMap: a non-zero timeEmbed produces a different output than the input (catches timestep-invariant outputs). Tests the helper methods directly rather than the full 5D forward pass — the full forward has pre-existing channel-tracking issues between SpatialResBlock (Dense on 4D image) and the decoder's concat-then-FiLM sequence that are outside the scope of this PR. Additional fixes discovered while writing tests: - ApplyFiLMConditioning now handles both 1D timeEmbed [channels*2] (broadcast across batches — the shared-across-batch case from GetTimestepEmbedding) and 2D timeEmbed [B, channels*2] (per-batch). Previously the validator treated condVec.Shape[0]=channels*2 as a batch mismatch. - ApplyTemporalProcessing and ApplyTemporalAttention now call .Contiguous() on their permute-then-reshape-then-permute-back results. Downstream ops (ExtractFrame, AsSpan callers) require contiguous backing buffers, and permutation views don't materialize automatically. - INeuralNetwork<T> : IDisposable (NeuralNetworkBase<T> already has Dispose, the interface just didn't expose it). Tests in ModelFamilyTests use `using var network = CreateNetwork()` which requires IDisposable on the declared interface type. Added Dispose() to MockNeuralNetwork and MockNeuralNetwork<T> in the test helpers. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * fix: PR #1137 round-5 — axis-correct temporal reshape + tests + FiLM 1D support Addresses 3 unresolved PR #1137 review comments: 1. ApplyTemporalProcessing: axis-correct temporal mixing. The reviewer correctly identified that a plain Reshape from [B, C, F, H, W] to [B*C*H*W, F] does NOT produce per-(b,c,h,w) temporal vectors — F is NOT the innermost dim in that layout, so the reshape enumerator mixes values across H/W into the "F slot". Fixed by permuting [B, C, F, H, W] → [B, C, H, W, F] BEFORE reshape to make F the innermost (row-major last) position, then permute back after the mixing. Mirrors the pattern ApplyTemporalAttention already uses (permute-before-reshape for the same axis-correctness reason). 2. CreateTemporalResBlock: removed misleading unused `channels` parameter. The layer is sized by `_numFrames`, not by channel count. Renamed to CreateTemporalMixingLayer() (per reviewer's suggested rename), removed the unused parameter, and updated all 4 call sites. The new name self-documents that it creates a frame-mixing layer. 3. Integration tests for timestep FiLM + temporal path. Added VideoUNetPredictorIntegrationTests exercising via reflection: - ApplyTemporalProcessing_PreservesShape: end-to-end [B,C,F,H,W] input produces matching-shape output (catches regression in the axis fix). - ApplyTemporalProcessing_MixesAcrossFramesAxis: per-frame-constant video stays per-frame-constant after mixing (catches the axis bug directly — a broken layout would mix H/W into the frame axis, producing non-constant-within-frame outputs). - ApplyFiLMConditioning_ModulatesFeatureMap: a non-zero timeEmbed produces a different output than the input (catches timestep-invariant outputs). Tests the helper methods directly rather than the full 5D forward pass — the full forward has pre-existing channel-tracking issues between SpatialResBlock (Dense on 4D image) and the decoder's concat-then-FiLM sequence that are outside the scope of this PR. Additional fixes discovered while writing tests: - ApplyFiLMConditioning now handles both 1D timeEmbed [channels*2] (broadcast across batches — the shared-across-batch case from GetTimestepEmbedding) and 2D timeEmbed [B, channels*2] (per-batch). Previously the validator treated condVec.Shape[0]=channels*2 as a batch mismatch. - ApplyTemporalProcessing and ApplyTemporalAttention now call .Contiguous() on their permute-then-reshape-then-permute-back results. Downstream ops (ExtractFrame, AsSpan callers) require contiguous backing buffers, and permutation views don't materialize automatically. - INeuralNetwork<T> : IDisposable (NeuralNetworkBase<T> already has Dispose, the interface just didn't expose it). Tests in ModelFamilyTests use `using var network = CreateNetwork()` which requires IDisposable on the declared interface type. Added Dispose() to MockNeuralNetwork and MockNeuralNetwork<T> in the test helpers. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * fix: address pr #1137 round-5 follow-up — 1d film support + contiguous materialization Applied on top of round-5 to cover edge cases the tests exposed: - ApplyFiLMConditioning now handles 1D [timeEmbedDim] timeEmbed by broadcasting the projection output across all batches. Before this change, the validator incorrectly flagged condVec.Shape[0] = channels*2 as a batch mismatch. Now detects condVec.Shape.Length < 2 as 1D and skips the batch check, reading from offset 0 for every batch index during the split. - ApplyTemporalProcessing and ApplyTemporalAttention now call .Contiguous() on their permute-then-reshape-then-permute-back results. Downstream ops that access raw spans (ExtractFrame in ProcessVideoFrames) fail on non-contiguous views, so these paths cannot return a permutation view. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * fix: PR #1137 round-5 follow-up — 1D FiLM support + Contiguous materialization Applied on top of round-5 to cover edge cases the tests exposed: - ApplyFiLMConditioning now handles 1D [timeEmbedDim] timeEmbed by broadcasting the projection output across all batches. Before this change, the validator incorrectly flagged condVec.Shape[0] = channels*2 as a batch mismatch. Now detects condVec.Shape.Length < 2 as 1D and skips the batch check, reading from offset 0 for every batch index during the split. - ApplyTemporalProcessing and ApplyTemporalAttention now call .Contiguous() on their permute-then-reshape-then-permute-back results. Downstream ops that access raw spans (ExtractFrame in ProcessVideoFrames) fail on non-contiguous views, so these paths cannot return a permutation view. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> --------- Co-authored-by: franklinic <franklin@ivorycloud.com> Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
1. Remove duplicate IAsyncLifetime methods in DiffusionModelTestBase (lines 88-104 were old versions; lines 42-63 are the LOH-compacting versions from master merge — kept those). 2. Add EnsureInitialized() to MultiHeadAttentionLayer.ForwardGpu() — lazy MHA on the GPU path would use [0,0] placeholder Q/K/V/O tensors, producing wrong outputs. 3. Add EnsureInitialized() to SelfAttentionLayer.ForwardGpu() — same lazy-init gap as MHA. 4. Add headCount validation to MultiHeadAttentionLayer constructor — headCount=0 now throws ArgumentOutOfRangeException instead of DivideByZeroException; non-divisible embedding size throws ArgumentException. 5. Merge duplicate Dispose loops in NeuralNetworkBase into a single DisposeOnceGuard-protected loop — prevents double-return of rented tensors to TensorAllocator. Build clean on net10.0. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Contributor
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
src/NeuralNetworks/Layers/MultiHeadAttentionLayer.cs (1)
1105-1113:⚠️ Potential issue | 🔴 CriticalMissing
EnsureWeightsAllocated()call—GPU path uses zero-sized placeholders.
ForwardGpu()only callsEnsureInitialized(), which handles sub-layer registration (ROPE/ALiBi) but NOT weight tensor materialization. A fresh lazy-initMultiHeadAttentionLayerentering the GPU path will project against[0,0]Q/K/V/O tensors, producing wrong outputs or crashes.The CPU path (
ForwardInternal) correctly calls both helpers.🔥 Proposed fix
public override Tensor<T> ForwardGpu(params Tensor<T>[] inputs) { if (inputs.Length == 0) throw new ArgumentException("At least one input tensor is required.", nameof(inputs)); // Materialize lazy Q/K/V/O weights before GPU projection matmuls. // Without this, a fresh lazy-init MHA entering the GPU path would // use [0,0] placeholder tensors and produce wrong outputs. + EnsureWeightsAllocated(); EnsureInitialized(); if (Engine is not DirectGpuTensorEngine gpuEngine)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/NeuralNetworks/Layers/MultiHeadAttentionLayer.cs` around lines 1105 - 1113, ForwardGpu currently calls EnsureInitialized() but misses materializing weight tensors, causing GPU path to use zero-sized Q/K/V/O placeholders; add a call to EnsureWeightsAllocated() immediately after EnsureInitialized() in the ForwardGpu method (same place ForwardInternal calls both) so weights are materialized before any GPU projection matmuls, ensuring Q/K/V/O tensors are valid for GPU execution.
♻️ Duplicate comments (4)
src/NeuralNetworks/Layers/MultiHeadAttentionLayer.cs (2)
357-387:⚠️ Potential issue | 🟠 MajorEager path ignores custom non-lazy
IInitializationStrategy<T>.When a caller passes a non-lazy custom strategy, the eager branch (lines 370-387) unconditionally calls
InitializeParameters(), applying default Xavier init. TheEnsureWeightsAllocated()helper correctly honors custom strategies, but it's never called in the eager path.🔧 Proposed fix
else { // Eager path: same behavior as before lazy init was added. _queryWeights = new Tensor<T>([embeddingDimension, embeddingDimension]); _keyWeights = new Tensor<T>([embeddingDimension, embeddingDimension]); _valueWeights = new Tensor<T>([embeddingDimension, embeddingDimension]); _outputWeights = new Tensor<T>([embeddingDimension, embeddingDimension]); _outputBias = new Tensor<T>([embeddingDimension]); - InitializeParameters(); + if (InitializationStrategy is { IsLazy: false }) + { + InitializationStrategy.InitializeWeights(_queryWeights, embeddingDimension, embeddingDimension); + InitializationStrategy.InitializeWeights(_keyWeights, embeddingDimension, embeddingDimension); + InitializationStrategy.InitializeWeights(_valueWeights, embeddingDimension, embeddingDimension); + InitializationStrategy.InitializeWeights(_outputWeights, embeddingDimension, embeddingDimension); + InitializationStrategy.InitializeBiases(_outputBias); + } + else + { + InitializeParameters(); + } RegisterTrainableParameter(_queryWeights, PersistentTensorRole.Weights);🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/NeuralNetworks/Layers/MultiHeadAttentionLayer.cs` around lines 357 - 387, The eager branch in MultiHeadAttentionLayer's constructor always calls InitializeParameters() and ignores a provided non-lazy IInitializationStrategy<T>; update the eager path to honor the supplied strategy by using the same allocation/initialization flow as the lazy case — call EnsureWeightsAllocated(initializationStrategy) (or otherwise pass the initializationStrategy into the allocation/initialization routine) to create and initialize _queryWeights, _keyWeights, _valueWeights, _outputWeights, and _outputBias with the custom strategy, then register them and set _isInitialized = true instead of unconditionally invoking the default InitializeParameters().
416-444:⚠️ Potential issue | 🟠 MajorSame issue: vector-activation constructor ignores custom non-lazy strategies.
Apply the same fix here as in the scalar-activation constructor.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/NeuralNetworks/Layers/MultiHeadAttentionLayer.cs` around lines 416 - 444, The vector-activation constructor currently ignores a non-lazy initializationStrategy and always calls InitializeParameters() with default behavior; update the else branch so it respects the supplied initializationStrategy: when initializationStrategy.IsLazy is false, construct the tensors (_queryWeights, _keyWeights, _valueWeights, _outputWeights, _outputBias) as shown and then call the initializer provided by initializationStrategy (e.g., pass initializationStrategy into InitializeParameters or invoke initializationStrategy.CreateInitializer()/Initialize(...) on each tensor) before RegisterTrainableParameter so the custom (non-lazy) strategy is applied rather than the default.src/NeuralNetworks/Layers/SelfAttentionLayer.cs (2)
387-421:⚠️ Potential issue | 🟠 MajorEager path still ignores custom non-lazy
IInitializationStrategy<T>.When a caller passes a non-lazy
initializationStrategy, the eager branch (lines 407-421) unconditionally callsInitializeParameters(), which applies the default Xavier init. The custom strategy is stored inInitializationStrategybut never invoked here.This means
new SelfAttentionLayer(..., initializationStrategy: myCustomEagerStrategy)silently ignoresmyCustomEagerStrategy. Only lazy strategies work as expected (sinceEnsureInitialized()eventually honors the strategy).🔧 Proposed fix
else { _queryWeights = new Tensor<T>([embeddingDimension, embeddingDimension]); _keyWeights = new Tensor<T>([embeddingDimension, embeddingDimension]); _valueWeights = new Tensor<T>([embeddingDimension, embeddingDimension]); _outputBias = new Tensor<T>([embeddingDimension]); - InitializeParameters(); + if (InitializationStrategy is { IsLazy: false }) + { + InitializationStrategy.InitializeWeights(_queryWeights, embeddingDimension, embeddingDimension); + InitializationStrategy.InitializeWeights(_keyWeights, embeddingDimension, embeddingDimension); + InitializationStrategy.InitializeWeights(_valueWeights, embeddingDimension, embeddingDimension); + InitializationStrategy.InitializeBiases(_outputBias); + } + else + { + InitializeParameters(); + } RegisterTrainableParameter(_queryWeights, PersistentTensorRole.Weights);🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/NeuralNetworks/Layers/SelfAttentionLayer.cs` around lines 387 - 421, The eager-construction branch ignores the provided InitializationStrategy and always calls InitializeParameters(), so update the constructor logic to use the configured InitializationStrategy (InitializationStrategy) when not lazy: after allocating _queryWeights, _keyWeights, _valueWeights, and _outputBias, invoke the strategy's initialization method (or a new helper that accepts an IInitializationStrategy<T>) to initialize those tensors instead of the hard-coded InitializeParameters(); keep RegisterTrainableParameter calls and _isInitialized=true, and ensure EnsureInitialized() still works for the lazy path.
472-502:⚠️ Potential issue | 🟠 MajorSame issue: vector-activation constructor ignores custom non-lazy strategies.
Apply the same fix here as in the scalar-activation constructor.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/NeuralNetworks/Layers/SelfAttentionLayer.cs` around lines 472 - 502, The vector-activation constructor currently only checks initializationStrategy.IsLazy and otherwise calls the default InitializeParameters, which ignores a provided custom non-lazy initializationStrategy; update the constructor so that after allocating _queryWeights, _keyWeights, _valueWeights and _outputBias you call the provided initializationStrategy to initialize those tensors when initializationStrategy != null (i.e. invoke the strategy's initialize method or equivalent) and fall back to InitializeParameters() only when no custom strategy is supplied, then RegisterTrainableParameter as before; reference symbols: initializationStrategy, _queryWeights, _keyWeights, _valueWeights, _outputBias, InitializeParameters, RegisterTrainableParameter.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Outside diff comments:
In `@src/NeuralNetworks/Layers/MultiHeadAttentionLayer.cs`:
- Around line 1105-1113: ForwardGpu currently calls EnsureInitialized() but
misses materializing weight tensors, causing GPU path to use zero-sized Q/K/V/O
placeholders; add a call to EnsureWeightsAllocated() immediately after
EnsureInitialized() in the ForwardGpu method (same place ForwardInternal calls
both) so weights are materialized before any GPU projection matmuls, ensuring
Q/K/V/O tensors are valid for GPU execution.
---
Duplicate comments:
In `@src/NeuralNetworks/Layers/MultiHeadAttentionLayer.cs`:
- Around line 357-387: The eager branch in MultiHeadAttentionLayer's constructor
always calls InitializeParameters() and ignores a provided non-lazy
IInitializationStrategy<T>; update the eager path to honor the supplied strategy
by using the same allocation/initialization flow as the lazy case — call
EnsureWeightsAllocated(initializationStrategy) (or otherwise pass the
initializationStrategy into the allocation/initialization routine) to create and
initialize _queryWeights, _keyWeights, _valueWeights, _outputWeights, and
_outputBias with the custom strategy, then register them and set _isInitialized
= true instead of unconditionally invoking the default InitializeParameters().
- Around line 416-444: The vector-activation constructor currently ignores a
non-lazy initializationStrategy and always calls InitializeParameters() with
default behavior; update the else branch so it respects the supplied
initializationStrategy: when initializationStrategy.IsLazy is false, construct
the tensors (_queryWeights, _keyWeights, _valueWeights, _outputWeights,
_outputBias) as shown and then call the initializer provided by
initializationStrategy (e.g., pass initializationStrategy into
InitializeParameters or invoke
initializationStrategy.CreateInitializer()/Initialize(...) on each tensor)
before RegisterTrainableParameter so the custom (non-lazy) strategy is applied
rather than the default.
In `@src/NeuralNetworks/Layers/SelfAttentionLayer.cs`:
- Around line 387-421: The eager-construction branch ignores the provided
InitializationStrategy and always calls InitializeParameters(), so update the
constructor logic to use the configured InitializationStrategy
(InitializationStrategy) when not lazy: after allocating _queryWeights,
_keyWeights, _valueWeights, and _outputBias, invoke the strategy's
initialization method (or a new helper that accepts an
IInitializationStrategy<T>) to initialize those tensors instead of the
hard-coded InitializeParameters(); keep RegisterTrainableParameter calls and
_isInitialized=true, and ensure EnsureInitialized() still works for the lazy
path.
- Around line 472-502: The vector-activation constructor currently only checks
initializationStrategy.IsLazy and otherwise calls the default
InitializeParameters, which ignores a provided custom non-lazy
initializationStrategy; update the constructor so that after allocating
_queryWeights, _keyWeights, _valueWeights and _outputBias you call the provided
initializationStrategy to initialize those tensors when initializationStrategy
!= null (i.e. invoke the strategy's initialize method or equivalent) and fall
back to InitializeParameters() only when no custom strategy is supplied, then
RegisterTrainableParameter as before; reference symbols: initializationStrategy,
_queryWeights, _keyWeights, _valueWeights, _outputBias, InitializeParameters,
RegisterTrainableParameter.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: ASSERTIVE
Plan: Pro
Run ID: 066a743b-4c8f-4306-b44a-d591e36c4f88
📒 Files selected for processing (4)
src/NeuralNetworks/Layers/MultiHeadAttentionLayer.cssrc/NeuralNetworks/Layers/SelfAttentionLayer.cssrc/NeuralNetworks/NeuralNetworkBase.cstests/AiDotNet.Tests/ModelFamilyTests/Base/DiffusionModelTestBase.cs
Four independent diffusion-noise-predictor conflicts, each analyzed
on semantic merit:
1. DiTNoisePredictor._labelEmbed output dim: took master's timeEmbedDim
(hiddenSize*4) over HEAD's _hiddenSize. Class embeddings are added to
the time-MLP output for AdaLN, so they must match the time-MLP output
space (timeEmbedDim), not the hidden_size. HEAD's version would cause
a runtime shape mismatch when adding class and time embeddings.
2. NoisePredictorBase: kept HEAD's LazyMHA + LazySelfAttention helper
methods — these are the core feature of this PR (defer attention
weight allocation until first Forward). Master doesn't have them.
3. VideoUNetPredictor (5 conflicts):
- Three were pure formatting (multi-line vs single-line args) — kept
HEAD's multi-line form.
- CreateTemporalMixingBlock: HEAD referenced undefined 'channels'
variable (compile error). Took master's _numFrames, which is the
semantically correct dimension for temporal-axis mixing.
- Also took master's new CreateTimeCondProjection helper (FiLM
conditioning projection, net-new in master).
4. NeuralNetworkModelTestBase: pure comment-text conflict, both sides
semantically identical (ParameterCount > 0 check). Kept HEAD's phrasing.
Build clean on net10.0.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
ooples
pushed a commit
that referenced
this pull request
Apr 16, 2026
Since this PR stacks on #1140, all the same conflicts recur plus a few new ones from the #1141-only changes. All resolved semantically: DiTNoisePredictor (3 conflicts): - EnsureLayersInitialized: deduplicated the _initLock field (master defines it once at line 182, HEAD redefined at line 312 — kept master's single definition + master's richer doc about retry-safety). - _labelEmbed output dim: took master's timeEmbedDim (correct for addition with time MLP output). Same semantic fix as PR #1140. NoisePredictorBase (1 conflict): kept HEAD's LazyMHA + LazySelfAttention helpers (core feature of the #1140 stack). VideoUNetPredictor (8 conflicts): - 4 were pure formatting (multi-line vs single-line args) — kept HEAD's multi-line form. - ApplyTemporalProcessing: took master's richer comments — same logic, but master documents WHY permute+reshape is needed (naive reshape would mix H/W instead of frames). - CreateTemporalMixingBlock: HEAD referenced undefined 'channels' (compile error). Took master's _numFrames + new CreateTimeCondProjection method. ConvolutionalLayer.Deserialize (1 conflict): took master's version — adds ClearRegisteredParameters() + RegisterTrainableParameter() calls so optimizers target the freshly-deserialized tensors. Without this, gradient updates silently go to stale references. NeuralNetworkModelTestBase (1 conflict): pure comment-text conflict, both semantically identical. Kept HEAD's phrasing. Build clean on net10.0. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
src/Diffusion/NoisePredictors/VideoUNetPredictor.cs (1)
291-316:⚠️ Potential issue | 🔴 CriticalDispose the privately-owned lazy layers from this model.
NeuralNetworkBase.Dispose(bool)only walksLayers, but this type keeps_inputConv,_outputConv,_imageCondProjection, and the block-owned layers in private fields/lists instead of registering them there. After this PR, those objects can hold pool-rented tensors, sousing var modelwill still leak the large buffers this change is trying to release.Suggested fix
+ protected override void Dispose(bool disposing) + { + if (disposing) + { + foreach (var layer in EnumerateOwnedLayers()) + { + if (layer is IDisposable disposable) + { + AiDotNet.Helpers.DisposeOnceGuard.TryDispose(disposable); + } + } + } + + base.Dispose(disposing); + } + + private IEnumerable<ILayer<T>> EnumerateOwnedLayers() + { + if (_inputConv is not null) yield return _inputConv; + if (_timeEmbedMlp1 is not null) yield return _timeEmbedMlp1; + if (_timeEmbedMlp2 is not null) yield return _timeEmbedMlp2; + if (_imageCondProjection is not null) yield return _imageCondProjection; + if (_outputConv is not null) yield return _outputConv; + + foreach (var block in _encoderBlocks.Concat(_middleBlocks).Concat(_decoderBlocks)) + { + if (block.SpatialResBlock is not null) yield return block.SpatialResBlock; + if (block.TemporalResBlock is not null) yield return block.TemporalResBlock; + if (block.SpatialAttention is not null) yield return block.SpatialAttention; + if (block.TemporalAttention is not null) yield return block.TemporalAttention; + if (block.CrossAttention is not null) yield return block.CrossAttention; + if (block.Downsample is not null) yield return block.Downsample; + if (block.Upsample is not null) yield return block.Upsample; + if (block.TimeCondProjection is not null) yield return block.TimeCondProjection; + } + }Also applies to: 328-359, 379-410
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/Diffusion/NoisePredictors/VideoUNetPredictor.cs` around lines 291 - 316, The model keeps privately-owned lazy layers (e.g., _inputConv, _outputConv, _imageCondProjection, _timeEmbedMlp1, _timeEmbedMlp2 and the block-owned layer lists) out of the base class Layers collection so they are not disposed; update Dispose(bool) in VideoUNetPredictor (or override Dispose) to explicitly Dispose() each of these private fields and any layers in the block-owned lists (safely checking for null and swallowing ObjectDisposed exceptions), or alternatively register them into the inherited Layers collection when created so NeuralNetworkBase.Dispose will clean them up; also set the fields to null after disposing to avoid double-dispose.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@src/Diffusion/NoisePredictors/VideoUNetPredictor.cs`:
- Around line 1034-1059: The spatial-attention and downsample factories are
incorrectly hardcoded to use _inputHeight/_inputWidth; update
CreateSpatialAttention and CreateDownsample (and any similar helpers) to accept
current per-level height and width parameters and use those when calling LazyMHA
and LazyConv2D instead of the global _inputHeight/_inputWidth, then thread the
current height/width through the encoder/decoder builder methods so each stage
calls CreateSpatialAttention(height, width, channels) and
CreateDownsample(height, width, channels) (or equivalent signatures) with the
progressively reduced spatial sizes; also update any callers of
CreateTemporalAttention/CreateCrossAttention only if they need analogous
per-level sizes.
In `@src/NeuralNetworks/Layers/ConvolutionalLayer.cs`:
- Around line 1492-1495: Dispose currently only returns rented _kernels when
_isInitialized is true, which leaks rentals if allocation occurs before
EnsureInitialized flips _isInitialized and then fails; update Dispose to return
rented kernels unconditionally (e.g., if _kernels != null && _kernels.Length >
0) rather than checking _isInitialized, call TensorAllocator.Return(_kernels)
and clear/null out _kernels after returning to avoid double-return, leaving
_isInitialized semantics unchanged and keeping EnsureInitialized() allocation
logic intact.
---
Outside diff comments:
In `@src/Diffusion/NoisePredictors/VideoUNetPredictor.cs`:
- Around line 291-316: The model keeps privately-owned lazy layers (e.g.,
_inputConv, _outputConv, _imageCondProjection, _timeEmbedMlp1, _timeEmbedMlp2
and the block-owned layer lists) out of the base class Layers collection so they
are not disposed; update Dispose(bool) in VideoUNetPredictor (or override
Dispose) to explicitly Dispose() each of these private fields and any layers in
the block-owned lists (safely checking for null and swallowing ObjectDisposed
exceptions), or alternatively register them into the inherited Layers collection
when created so NeuralNetworkBase.Dispose will clean them up; also set the
fields to null after disposing to avoid double-dispose.
🪄 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: 705a11b7-47aa-4de0-b3ef-c6fd601e74ab
📒 Files selected for processing (4)
src/Diffusion/NoisePredictors/DiTNoisePredictor.cssrc/Diffusion/NoisePredictors/VideoUNetPredictor.cssrc/NeuralNetworks/Layers/ConvolutionalLayer.cssrc/NeuralNetworks/NeuralNetworkBase.cs
1. ConvolutionalLayer.Dispose: changed Dispose gate from `_isInitialized && _kernels.Length > 0` to just `_kernels.Length > 0`. EnsureInitialized rents the tensor BEFORE flipping _isInitialized to true. If weight population throws between those steps, Dispose would leak the rented tensor because _isInitialized is still false. Length > 0 is a sufficient check since lazy placeholders sit at Length == 0 and aren't pool-rented. 2. VideoUNetPredictor: per-level layer factories now take a 'level' parameter and compute resolution as _inputHeight >> level. The old code baked _inputHeight/_inputWidth (top-level resolution) into every block's SpatialAttention/Downsample/Upsample — correct only for level 0, wrong for deeper stages where feature maps are 2x smaller per level. With lazy init, the wrong sizing now materializes at first Forward. Added ResolutionAtLevel(level) helper. Updated callers in InitializeLayers: encoder passes its current level, middle passes the deepest level (channelMultipliers.Length - 1), decoder passes its level. CreateUpsample takes decoder's level and computes input resolution as level+1 (halved from the paired encoder level). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
ooples
pushed a commit
that referenced
this pull request
Apr 16, 2026
1. VideoUNetPredictor cross-attention fallback: changed from silent
self-attention fallback to throwing InvalidOperationException when
CrossAttention isn't a LayerBase<T>. Silently dropping the text
conditioning tensor and degrading to self-attention is a correctness
regression that'd be extremely hard to debug from outputs alone.
2. ConvolutionalLayer.Dispose:
- Changed `_isInitialized && _kernels.Length > 0` gate to just
`_kernels.Length > 0` (same fix as PR #1140 — partial init failure
leaves Length > 0 but flag still false, leaking the rented tensor).
- NEW: also return the rented _preAllocatedOutput forward buffer.
Without this, disposing many ConvolutionalLayer instances (one
per conv in a deep UNet) leaks one rented activation per layer —
dozens of MB per disposed model at SD scale.
3. VideoUNet per-level resolution fix (same as PR #1140): factories
now take a 'level' parameter and compute resolution as
_inputHeight >> level instead of baking the top-level resolution
into every block. Required because PR #1141 stacks on #1140.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Contributor
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
src/Diffusion/NoisePredictors/VideoUNetPredictor.cs (2)
1075-1089:⚠️ Potential issue | 🔴 CriticalCritical: Off-by-one error in upsample input resolution calculation.
The
CreateUpsamplemethod usesResolutionAtLevel(level + 1)but this is incorrect. When the decoder is at level N and adds an upsample block, the input to that upsample comes from processing at level N's resolution, which isres(N)=_inputHeight >> level, NOT_inputHeight >> (level + 1).Concrete example (4 levels,
_inputHeight=64):
- Encoder downsamples: 64→32→16→8 (deepest =
res(3)= 8)- Decoder level 3 upsample: input is 8×8 (
res(3)), output should be 16×16- Code computes:
ResolutionAtLevel(4)=64 >> 4= 4 ❌- Correct:
ResolutionAtLevel(3)=64 >> 3= 8 ✓This will cause a shape mismatch at runtime when
DeconvolutionalLayerreceives input at the wrong expected dimensions.🐛 Proposed fix
private ILayer<T> CreateUpsample(int channels, int level) { - // Decoder level N has input resolution _inputHeight >> (level+1) - // (output of the corresponding encoder downsample) and upsamples to - // _inputHeight >> level (the paired encoder-level resolution). - int inputRes = ResolutionAtLevel(level + 1); + // Decoder level N receives input at resolution _inputHeight >> level + // (matching the encoder's resolution at that level) and upsamples to + // _inputHeight >> (level - 1), the next coarser encoder resolution. + int inputRes = ResolutionAtLevel(level); // Transposed convolution: stride=2, kernel=4, padding=1 ⇒ output = 2 * input return new DeconvolutionalLayer<T>( inputShape: new[] { 1, channels, inputRes, inputRes },🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/Diffusion/NoisePredictors/VideoUNetPredictor.cs` around lines 1075 - 1089, The upsample input resolution is computed incorrectly in CreateUpsample: it uses ResolutionAtLevel(level + 1) but the decoder's upsample receives input at res(level) (e.g., decoder level N input is _inputHeight >> level). Fix CreateUpsample by calling ResolutionAtLevel(level) for inputRes so the DeconvolutionalLayer<T> is constructed with the correct inputShape and avoids runtime shape mismatches.
1262-1264:⚠️ Potential issue | 🟠 MajorBlocking: Empty
BackwardBlockmethod violates production-readiness guidelines.Empty method bodies are explicitly flagged as production-readiness violations. Either:
- Remove this as dead code if it's not called
- Implement the method with actual logic
- If intentionally a no-op (gradient computation handled by GradientTape elsewhere), add an explanatory comment
private static void BackwardBlock(VideoBlock block, ref Tensor<T> grad) { + // No-op: Gradient computation is handled by GradientTape autodiff. + // Manual backward pass through individual blocks is not required. }Or delete entirely if unused. Do not leave empty method bodies in production code.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/Diffusion/NoisePredictors/VideoUNetPredictor.cs` around lines 1262 - 1264, The BackwardBlock(VideoBlock block, ref Tensor<T> grad) method is empty and must be fixed: either implement the backward/gradient propagation logic for VideoBlock (ensure gradients are accumulated into the ref Tensor<T> grad matching the forward pass in the corresponding ForwardBlock), remove the method entirely if it is unused, or if it's intentionally a no-op because gradients are handled elsewhere (e.g., by a GradientTape mechanism), add a clear explanatory comment above BackwardBlock stating that it's intentionally empty and why; update any callers or unit tests accordingly to avoid silent regressions.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Outside diff comments:
In `@src/Diffusion/NoisePredictors/VideoUNetPredictor.cs`:
- Around line 1075-1089: The upsample input resolution is computed incorrectly
in CreateUpsample: it uses ResolutionAtLevel(level + 1) but the decoder's
upsample receives input at res(level) (e.g., decoder level N input is
_inputHeight >> level). Fix CreateUpsample by calling ResolutionAtLevel(level)
for inputRes so the DeconvolutionalLayer<T> is constructed with the correct
inputShape and avoids runtime shape mismatches.
- Around line 1262-1264: The BackwardBlock(VideoBlock block, ref Tensor<T> grad)
method is empty and must be fixed: either implement the backward/gradient
propagation logic for VideoBlock (ensure gradients are accumulated into the ref
Tensor<T> grad matching the forward pass in the corresponding ForwardBlock),
remove the method entirely if it is unused, or if it's intentionally a no-op
because gradients are handled elsewhere (e.g., by a GradientTape mechanism), add
a clear explanatory comment above BackwardBlock stating that it's intentionally
empty and why; update any callers or unit tests accordingly to avoid silent
regressions.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: ASSERTIVE
Plan: Pro
Run ID: 5de019fe-e863-4252-823d-a14197bd904e
📒 Files selected for processing (2)
src/Diffusion/NoisePredictors/VideoUNetPredictor.cssrc/NeuralNetworks/Layers/ConvolutionalLayer.cs
ooples
pushed a commit
that referenced
this pull request
Apr 17, 2026
…s (PR #1141) Master merged PR #1140 so stacked conflicts surfaced in 4 files: DiffusionModelTestBase.cs: removed duplicate InitializeAsync + DisposeAsync methods (CS0111 compile error). Master's LOH-compacting pair at lines 42/63 is kept; the older plain-GC pair at 88/98 was dead code from a prior merge resolution that forgot to delete the obsolete pair. VideoUNetPredictor.cs CreateUpsample: trivial comment-ordering conflict — kept consolidated comment that also documents the ResolutionAtLevel square-input assumption (addresses review nitpick about non-square inputs producing wrong attention sequence lengths). ConvolutionalLayer.cs: combined both sides' improvements — master's richer Length > 0 rationale comment + this PR's new _preAllocatedOutput pool-return fix (master doesn't have that yet). SelfAttentionLayer.cs: kept HEAD's headCount validation guards (positive check + divisibility check with named params) in BOTH constructors. Master's weaker single-line ArgumentException is dropped. The guards run BEFORE the division so headCount=0 throws ArgumentOutOfRangeException cleanly instead of DivideByZeroException. Build clean on net10.0 for both src/AiDotNet.csproj and tests/AiDotNet.Tests. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
ooples
added a commit
that referenced
this pull request
Apr 17, 2026
…1141) * perf: lazy weight init in diffusion noise predictors (part 1/5 of #1136) Every diffusion model test on CI was OOMing during construction because DiT/MMDiT/UNet noise predictors eagerly allocate ~4 GB of weight tensors in their ctors — `new DenseLayer<T>(...)` calls `TensorAllocator.Rent<T>` before the model has even seen an input. With 255 diffusion models × ~10 tests each running in a shared xunit process, sequential tests stack up and OOM on the 16 GB Windows CI runners, cancelling all three Diffusion ModelFamily shards at the 45-minute wall clock. This part of the fix makes model construction O(1) by threading `InitializationStrategies<T>.Lazy` into every internal Dense/Conv layer built by a noise predictor. Weight tensors stay at shape [0,0] until the first Forward() pass actually needs them. Adds two protected helpers on NoisePredictorBase<T>: - LazyDense(int, int, IActivationFunction?) — lazy dense layer - LazyDenseVec(int, int, IVectorActivationFunction) — lazy dense for vector activations (distinct name avoids ctor-overload ambiguity on activations that implement both interfaces) - LazyConv2D(...) — lazy 2D convolutional layer Converts all 11 noise predictors + DiffusionResBlock: - DiTNoisePredictor (patch/time embeds + final + MLP/attention blocks) - MMDiTNoisePredictor (image/text streams + joint + single blocks) - MMDiTXNoisePredictor - EMMDiTPredictor - FlagDiTPredictor - FluxDoubleStreamPredictor (double + single streams) - AsymmDiTPredictor - SiTPredictor - UViTNoisePredictor (encoder + decoder + skip projections) - UNetNoisePredictor (downsample) - VideoUNetPredictor (time embed + spatial/temporal res blocks + in/out conv) - DiffusionResBlock (norm/conv1 + timeMlp + norm/conv2 + skip conv) Parts 2-5 (MHA/LayerNorm lazy init, Dispose→pool return, test lifecycle, GetParameters→ParameterCount swap) follow in subsequent commits. Refs #1136 * perf: lazy-friendly Parameters_ShouldBeNonEmpty + GC between tests (parts 4+5/5 of #1136) Two follow-on pieces to the lazy-init-in-noise-predictors change (PR 1 of the #1136 series). Together with that first commit they should clear the three cancelled Diffusion ModelFamily shards on CI — if not fully, they put the next diagnostics in reach instead of hiding behind OOM. Part 5: swap `GetParameters().Length > 0` for `ParameterCount > 0`. The `Parameters_ShouldBeNonEmpty` tests in DiffusionModelTestBase and NeuralNetworkModelTestBase are just asking "does this model have any learnable parameters?" — a question `ParameterCount` answers in O(1) without materializing the full flattened parameter vector. Calling GetParameters() on a lazily-constructed DiT-XL forces every DenseLayer to eagerly allocate weights (EnsureInitialized cascades down) just to count them — that alone reproduces the ~4 GB OOM even with lazy init in place. ParameterCount propagates through the layer list without triggering lazy materialization, so the existence check stays cheap. Semantically identical assertion, not a weakening — both check that the model has learnable parameters. Just using the right API. Part 4: IAsyncLifetime.DisposeAsync forces a full GC cycle between test methods in both ModelFamily bases. Each test instantiates a fresh production-sized model (VGG16BN, DiT-XL, SDXL, etc.). Without explicit GC pressure, the shared xunit process holds onto previous test's weight tensors until the collector runs on its own schedule. Sequential tests in a shard stack up live weight allocations that exceed the 16 GB Windows runner budget. IAsyncLifetime.DisposeAsync runs after every [Fact] in a class, so a Collect/WaitForPendingFinalizers/Collect cycle there releases the previous test's weight tensors before the next test's model constructs. Combined with Part 1 (lazy construction) and Part 5 (no forced materialization on existence check), the working-set for a single test stays bounded to what that test actually touches. Parts 2 (MHA/LayerNorm lazy init) and 3 (Dispose → pool-return) are deferred to a follow-up PR once we see how much of the OOM is addressed by Parts 1/4/5 alone. Refs #1136 * perf: lazy attention weights + Dispose cascades pool return (parts 2+3/5 of #1136) Second round of memory fixes for the cancelled Diffusion / NeuralNetworks CI shards. Parts 1+4+5 in PR #1137 landed lazy DenseLayer/ConvolutionalLayer init in the noise predictors but didn't move the needle — Diffusion shards still cancelled at 45 min. The diagnosis: each DiT-XL tower still eagerly allocates ~1 GB of attention weights across 28 blocks' worth of Q/K/V projections, because neither `SelfAttentionLayer` nor `MultiHeadAttentionLayer` had a lazy-init path. This PR adds it. ## Part 2 — Lazy init for MHA + SelfAttentionLayer Both layers gain an optional `IInitializationStrategy<T>?` constructor parameter. When the strategy is `InitializationStrategies<T>.Lazy`, Q/K/V/O weight tensors stay at shape `[0,0]` until the first Forward() call: - `SelfAttentionLayer`: straightforward — `EnsureInitialized()` override allocates on first Forward / GetParameters / SetParameters. No sub-layer fields, so no TrainableParameterGenerator conflict. - `MultiHeadAttentionLayer`: trickier. The generator already emits an `EnsureInitialized` override on MHA because `_ropeLayer` and `_alibiLayer` are sub-layer fields needing registration. Our lazy-init logic lives in a separate private helper `EnsureWeightsAllocated()` called from Forward / GetParameters / SetParameters. The generator-emitted EnsureInitialized still runs its sub-layer registration. Both expose lazy-friendly `ParameterCount` so existence checks don't force allocation. `LazyMHA` and `LazySelfAttention` factories added to `NoisePredictorBase<T>` alongside the existing `LazyDense` / `LazyConv2D` helpers. `DiTNoisePredictor.CreateAttentionLayer`, `UViTNoisePredictor`, and `VideoUNetPredictor.CreateSpatial/Temporal/CrossAttention` updated. Skipped `LayerNormalizationLayer` — gamma/beta are `[featureSize]` tensors (~18 KB each), not a meaningful OOM contributor. ## Part 3 — Dispose cascades rented tensors back to the pool - `ConvolutionalLayer.Dispose` now calls `TensorAllocator.Return(_kernels)` in the eager-init branch, matching the pattern `DenseLayer.Dispose` already has for `_weights` / `_biases`. The convolutional kernels are rented via `RentUninitialized` in the eager path, so they need to go back to the pool for reuse by the next model. - `NeuralNetworkBase.Dispose(bool)` now cascades to every layer that implements `IDisposable`, catching `ObjectDisposedException` so a cross-network layer share doesn't abort the loop. - `DiffusionModelBase` gains its own `Dispose` / `Dispose(bool)` since it doesn't inherit from `NeuralNetworkBase`. Subclasses can override to cascade to their own disposables. - `INeuralNetworkModel<T>` and `IDiffusionModel<T>` now inherit `IDisposable` so callers can wrap models in `using var`. `NeuralNetworkBase` and `DiffusionModelBase` already implement the contract — this just tightens the interface surface to match what concrete types already do. - Test bases `DiffusionModelTestBase` and `NeuralNetworkModelTestBase` now wrap `CreateModel()` / `CreateNetwork()` results in `using var` for all 29 test methods. Combined with the `IAsyncLifetime.DisposeAsync` GC hook from PR #1137, each test now actually releases its weight buffers back to the allocator pool before the next test constructs. ## Stacked on #1137 Branched off `perf/diffusion-lazy-init-oom`. Merges naturally either way (this PR's diff against master after #1137 lands cleanly covers only the Parts 2+3 additions). Refs #1136 * perf: skip pre-swap parameter walk in TrainWithTape steady state (#1136) TrainWithTape was doing a full recursive CollectParameters walk every Train() call before the buffer swap, just to size the parameter buffer. The buffer only needs sizing on the first Train() call (when it doesn't exist). Every iteration after that, the buffer is already correctly sized — GetOrCreateParameterBuffer short-circuits and returns the existing one without needing the sizing input. Skip the pre-swap walk when the buffer already exists. The post-swap walk still runs (it's the authoritative one, returning buffer-view references that match the tape). On DiT-XL with 28 transformer blocks (each recursively containing MLP + attention + norms + cross-attn projections) this saves one full layer-tree traversal + parameter deduplication per Train() call. Scope-limited slice of the runtime perf work tracked in #1136. Broader scalar-loop-to-Engine-op conversions (GraphTransformerLayer's manual QK^T, MessagePassingLayer's per-edge message computation, SSM variants) are NOT in the hot path for the currently-cancelled Diffusion / NeuralNetworks shards and don't move the CI needle. The actual perf bottleneck for those shards lives in the external AiDotNet.Tensors package (TensorMatMul, Conv2D, ScaledDotProductAttention on CPU) — tracked separately. Refs #1136 * perf: replace scalar NumOps loops with Engine ops in two more layers (#1136) Continues the runtime-perf work from the previous commit. Replaces two scalar-loop hot paths with fused Engine operations. Each converted loop was dispatching per-element virtual NumOps.Add/Multiply/Subtract calls, which JIT cannot vectorize through the IInterface<T> boundary; Engine ops dispatch once and run the whole tensor through a SIMD/GPU- accelerated kernel. IntersampleAttentionLayer.AddResidualAndNormalize: Manual batch-normalization across the embedding dimension was doing 3 × batchSize × numFeatures × embDim scalar NumOps calls (one pass for mean, one for variance, one for normalize+gamma/beta). Replaced with a single Engine.LayerNorm call with identical semantics. DepthwiseSeparableConvolutionalLayer.PointwiseConvolution: A 1×1 convolution is a linear projection over the channel dimension, equivalent to matmul. The previous nested-loop implementation was O(B·H·W·OC·IC) scalar NumOps.Multiply calls — for a MobileNet block at 224×224 input with 32→64 channels that's ~100M virtual dispatches per forward. Replaced with Engine.Reshape + Engine.TensorMatMul: flatten to [B·H·W, IC] @ [IC, OC], reshape back to [B, H, W, OC]. The kernel transpose goes through Engine.TensorPermute so no scalar loop remains. Refs #1136 * perf: replace BatchEnsembleLayer member-averaging scalar loop with ReduceMean AverageMembers was a classic nested-loop mean-over-axis: for every batch item × output column, sum across numMembers and divide. Each iteration did one NumOps.Add dispatch and the outer divide did NumOps.Multiply — per-forward cost O(batchSize × outputDim × numMembers) virtual calls. Reshape [B·M, OutputDim] → [B, M, OutputDim] and reduce-mean over the member axis is semantically identical and runs as one Engine op. Refs #1136 * perf: replace manual LayerNorm/GroupNorm forward loops in SSM layers (#1136) Six SSM-family layers had hand-written normalization implementations doing 4 nested scalar passes per (batch, time) position (compute mean, compute variance, compute invStd, apply normalize+scale+bias). Each pass is batchSize × seqLen × {numHeads ×} headDimension virtual NumOps dispatches that JIT cannot vectorize. All six are equivalent to either Engine.LayerNorm (when normalizing across the entire feature axis) or Engine.GroupNorm (when normalizing per-head independently with a shared mean/variance). Conversions: - RWKV7Block.ApplyGroupNorm — per-head GroupNorm over modelDim - LonghornLayer.ApplyGroupNorm — same pattern, [B, T, D] input - RetNetLayer.GroupNormForward — captures mean/variance for backward - TTTLayer.LayerNormForward — straight LayerNorm over last axis - MesaNetLayer.LayerNormForward — straight LayerNorm over last axis - MegalodonLayer (timestep norm) — straight LayerNorm; stdInv derived from variance via vectorized Engine ops (TensorAddScalar + TensorSqrt + TensorReciprocal) for backward Backward passes for these layers retain their manual implementations — they have hand-derived gradient chains that need an Engine.LayerNormBackward follow-up to fully fuse. Forward is the dominant call (one backward per training iteration vs. potentially many forwards in inference / iterative loops), so these are the higher-impact half. Refs #1136 * perf: GroupedQueryAttention SDPA + GatedFeatureLearning ReduceMean (#1136) Two more scalar-loop replacements: - GroupedQueryAttentionLayer.ComputeStandardAttention: the manual Q·K^T → softmax → attn·V was 6 nested loops doing per-element NumOps dispatches (O(B·H·SeqQ·SeqKV·D) per Q·K^T pass). Replaced with one Engine.ScaledDotProductAttention call which fuses scale + softmax + attn·V into a single SIMD/GPU kernel and returns the attention weights via the existing out-param contract. - GatedFeatureLearningUnitLayer.GetFeatureImportance: per-output-dim scalar accumulation across batch axis replaced with one Engine.ReduceMean over axis 0. Refs #1136 * perf: replace TabMRegression scalar reductions and loss loops (#1136) TabMRegression had three scalar-loop hotspots in its prediction and loss APIs that fired per-batch: - PredictWithUncertainty: nested per-batch × output-dim loops computing ensemble mean and variance via scalar NumOps.Add/Multiply/Subtract. Replaced with a [B, M, OutputDim] reshape + Engine.ReduceMean for the mean, broadcast-subtract + square + ReduceMean for the variance, and Engine.TensorSqrt for the std-dev. The per-batch O(M × OutputDim) scalar dispatch chain collapses to 5 Engine calls regardless of batch or member count. - ComputeMSELoss: per-element scalar (predicted - target)² accumulation replaced with TensorSubtract + TensorMultiply + TensorSum. - ComputeMAELoss: per-element scalar |predicted - target| accumulation replaced with TensorSubtract + TensorAbs + TensorSum. Refs #1136 * perf: vectorize DepthwiseConv per-channel + GraphTransformer attention (#1136) Two patterns I'd previously claimed were too complex to convert. Both are doable with the Engine ops already exposed; PyTorch handles them the same way. DepthwiseSeparableConvolutionalLayer.DepthwiseConvolution: Each input channel is convolved with its own [KH, KW] kernel independently — that's the depthwise op. The previous nested-loop scalar implementation was O(B·H·W·IC·KH·KW) per forward pass. Engine.Conv2D doesn't expose a `groups` parameter (filed as AiDotNet.Tensors#162), but the same effect is achievable by permuting NHWC→NCHW once, slicing per channel, calling Engine.Conv2D on each [B, 1, H, W] slice with the corresponding [1, 1, KH, KW] kernel, and concatenating along the channel axis before permuting back. Per-channel slice/concat overhead is bounded by IC iterations (32 for typical MobileNet) — small relative to the SIMD/GPU win on the convolution kernel itself. GraphTransformerLayer.MultiHeadAttention: The manual quad-nested loop computing Q·K^T element-by-element, then softmax, then attn·V was the textbook scaled dot-product attention with two graph-specific extras (additive structural bias + adjacency mask). Replaced with: keysT = TensorPermute(K, [0, 2, 1]) scores = TensorMatMul(Q, keysT) * (1/sqrt(d_k)) scores += sliced structural bias [1, N, N] broadcast over batch scores += adjacency mask additive form: (adj * 1e9 - 1e9) precomputed once per call, broadcast over batch attnWeights = Softmax(scores) headOut = TensorMatMul(attnWeights, V) The adjacency-mask trick is the standard PyTorch pattern: encode "where adj=0 set score=-inf" as an additive mask of (-1e9) at those positions; softmax drives the corresponding weights to ~0. Avoids branching inside the inner loop. Per-head Q/K/V/attn/output scatter into the 4D caches still uses an explicit write loop (no Engine "set-slice" op), but that's bounded to O(B·N·D) per head — vs the O(B·H·N²·D) compute that the matmul rewrite collapsed. Refs #1136 * perf: vectorize MessagePassingLayer end-to-end (#1136) The whole forward pass was scalar. PyTorch / torch_geometric handles this exact pattern with bulk tensor ops; replicating that here: Step 1+2 (message MLP + masked aggregation): - Per-edge messageInput build: replaced quintuple-nested loop with Engine.TensorTile broadcasts of node features into [B, N, N, F] tiles (src indexed by j, tgt indexed by i) plus optional edge feature reshape, all concatenated along the feature axis. - 2-layer MLP: flatten [B, N, N, MessageInputDim] → [B·N·N, ...], run two TensorMatMul + bias + ReLU passes vectorized over the whole edge set instead of one (b,i,j,h,k) chain at a time. - Adjacency masking: multiply messages by adj broadcast across the feature axis. Equivalent to the original `if adj==0 skip` since zeros pre-aggregation contribute nothing to the sum. - Aggregation: ReduceSum over the j axis instead of nested per-(b,i,h) scalar accumulation. Step 3 (GRU update): - Reset and update gates: now σ(input @ W + agg @ Wm + b) on the flattened [B·N, *] input via TensorMatMul + TensorBroadcastAdd + Sigmoid. Same gate semantics, two matmuls per gate instead of per-cell scalar accumulation across input + message features. - Output combiner: out = (1 - update) * pad(input) + update * pad(agg). The original conditional `f < inputFeatures ? input[f] : 0` is expressed as a slice/pad of the last axis to outputFeatures (PadOrSliceLastAxis helper). Lets the elementwise multiply/add use bulk Engine ops on equal-width tensors. Refs #1136 * perf: lazy EmbeddingLayer initialization (#1136) EmbeddingLayer's tensor was eagerly allocated in the constructor — [vocabularySize, embeddingDimension]. For BERT-scale transformers (BGE / SGPT / Matryoshka), that's ~30,522 × 768 × 8 bytes ≈ 187 MB materialized at `new EmbeddingLayer(...)` time, before any test input flows through. Multiple EmbeddingLayer instances per network × multiple test methods × shared xunit process = hundreds of MB held live across test boundaries. Made initialization lazy following the same pattern PR #1140 introduced for MultiHeadAttentionLayer / SelfAttentionLayer: - Cached _vocabularySize / _embeddingDimension / _embeddingInitialized fields. ParameterCount reads from the cached fields, not from the placeholder tensor's [0,0] shape — stays correct without forcing materialization. - Constructor leaves _embeddingTensor at [0,0] until first access. EnsureEmbeddingInitialized() allocates the real shape, runs the same SimdRandom-scaled fill the constructor used to do, registers with the engine for GPU persistence. - Wired into Forward, GetParameters, GetTokenEmbeddings, SetParameters so any data-touching path materializes first. SetParameters also short-circuits the lazy flag since it's writing the real-sized tensor itself. Refs #1136 * perf: vectorize BatchEnsembleLayer.Forward (#1136) The batch-ensemble Forward had two scalar-loop hotspots: - Per-(b, m, i) build of expandedInput / scaledInput by tiling input M times and elementwise-multiplying by per-member r-vectors. - Per-(row, j, i) ensemble matmul + per-row s-vector scaling + bias. Both replaced with batched Engine ops: - Tile input via Reshape → TensorTile → Reshape (collapses (b, m) into the row axis matching the original [b*M + m, i] layout). Same pattern applied to _rVectors and _sVectors so per-row scaling is one TensorMultiply instead of a per-element scalar dispatch. - Matmul via Engine.TensorMatMul on the [B*M, inputDim] @ [inputDim, outputDim] reshape of weights. - Bias add via TensorBroadcastAdd of the shared [outputDim] bias. Refs #1136 * perf: GraphAttentionLayer attention-weight slice via Engine.TensorSlice (#1136) The sparse-aggregation and dense-aggregation paths both extracted the source / target halves of _attentionWeights via per-(h, f) scalar copy loops to build separate tensors. Replaced with Engine.TensorSlice on the [numHeads, 2*outputFeatures] field — slices columns [0, outputFeatures) for source and [outputFeatures, 2*outputFeatures) for target as view-style operations in one call instead of building scratch tensors element by element. Refs #1136 * fix: address 10 review comments with actual code fixes (PR #1141) Real code fixes, not just thread resolution: 1. DiTNoisePredictor.PredictNoiseWithEmbedding: added missing EnsureLayersInitialized() call — fresh lazy instance would throw InvalidOperationException on first PredictNoiseWithEmbedding. 2. ConvolutionalLayer.Deserialize: set _isInitialized = true after renting _kernels via TensorAllocator — without this, Dispose skipped the Return call and deserialized layers leaked their rented tensors. 3. EmbeddingLayer.ForwardGpu: added EnsureEmbeddingInitialized() — lazy embedding on GPU path would use [0,0] placeholder and produce wrong lookups. 4. EmbeddingLayer.GetMetadata: use _vocabularySize/_embeddingDimension stored fields instead of _embeddingTensor.Shape — lazy placeholder [0,0] would serialize as VocabularySize=0, EmbeddingDimension=0. 5. MessagePassingLayer: return rented zeros tensor to TensorAllocator after TensorConcatenate — previously leaked every padding call. 6. SelfAttentionLayer: move headCount validation BEFORE the division (both constructors) — headCount=0 now throws ArgumentOutOfRangeException instead of DivideByZeroException, and non-divisible embedding sizes throw before the integer truncation occurs. 7. GraphTransformerLayer: remove unused maskNeg variable (dead code). Build clean on net10.0. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * fix: fix 3 remaining pre-existing bugs found by reviewer (PR #1141) 1. VideoUNet ApplyTemporalProcessing: was a no-op clone that never called the temporal layer. Replaced with correct permute→reshape→Forward→ reshape→permute + residual connection. The temporal mixing layer now actually processes the frame axis. 2. VideoUNet cross-attention: Forward(frame) was self-attention (Q=K=V=frame), not cross-attention (Q=frame, K=V=conditioning). Cast CrossAttention to LayerBase<T> to access Forward(query, kv) overload. Conditioning tensor now flows through to the attention mechanism as intended. 3. TabMRegression loss helpers: improved error messages in ComputeMSELoss and ComputeMAELoss to include both length and shape in the exception message, making shape-mismatch debugging easier. Build clean on net10.0. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * fix: thread-safe DiT lazy init via double-checked locking (PR #1141) EnsureLayersInitialized in DiTNoisePredictor was a plain if-check with no synchronization. Two concurrent PredictNoise calls on the same model instance (e.g., request-pool serving) could race: both see _layersInitialized=false, both enter InitializeLayers, and the model ends up with a double set of blocks. Added the double-checked locking pattern (volatile flag + lock + re-check inside lock) matching the thread-safety PR #1137 already uses on its branch. The _blocks.Clear() inside the lock makes retries after partial init failure safe. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * fix: 2 real bugs found by review on PR #1141 1. VideoUNetPredictor cross-attention fallback: changed from silent self-attention fallback to throwing InvalidOperationException when CrossAttention isn't a LayerBase<T>. Silently dropping the text conditioning tensor and degrading to self-attention is a correctness regression that'd be extremely hard to debug from outputs alone. 2. ConvolutionalLayer.Dispose: - Changed `_isInitialized && _kernels.Length > 0` gate to just `_kernels.Length > 0` (same fix as PR #1140 — partial init failure leaves Length > 0 but flag still false, leaking the rented tensor). - NEW: also return the rented _preAllocatedOutput forward buffer. Without this, disposing many ConvolutionalLayer instances (one per conv in a deep UNet) leaks one rented activation per layer — dozens of MB per disposed model at SD scale. 3. VideoUNet per-level resolution fix (same as PR #1140): factories now take a 'level' parameter and compute resolution as _inputHeight >> level instead of baking the top-level resolution into every block. Required because PR #1141 stacks on #1140. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> --------- Co-authored-by: franklinic <franklin@ivorycloud.com> Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
6 tasks done
This was referenced Apr 27, 2026
This was referenced May 17, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Stacks on PR #1137. Parts 2+3 of the issue #1136 fix plan — the pieces PR #1137 deferred because they touched core layer infrastructure. After seeing that PR #1137 alone moved only 1 of 6 cancelled Diffusion/NN shards to SUCCESS, the follow-up work is essential: each DiT-XL tower still eagerly allocates ~1 GB of attention weights across 28 blocks' worth of Q/K/V/O projections, because
MultiHeadAttentionLayerandSelfAttentionLayerhad no lazy-init path.Part 2 — Lazy attention weights
MultiHeadAttentionLayerandSelfAttentionLayergain an optionalIInitializationStrategy<T>?ctor parameter. WhenIsLazy, Q/K/V/O weight tensors stay at shape[0,0]until first Forward() / GetParameters() / SetParameters() call.SelfAttentionLayer: standardEnsureInitialized()override. No sub-layer fields → no generator conflict.MultiHeadAttentionLayer: theTrainableParameterGeneratoralready emits anEnsureInitializedoverride on MHA because of its_ropeLayer/_alibiLayersub-layer fields. Our lazy-init logic lives in a separate privateEnsureWeightsAllocated()helper called from Forward/GetParameters/SetParameters. The generator's own override still runs sub-layer registration.ParameterCountcomputed from the stored_embeddingDimensionfield rather than forcing tensor allocation for existence checks.NoisePredictorBase<T>gainsLazyMHAandLazySelfAttentionfactory helpers alongside the existingLazyDense/LazyDenseVec/LazyConv2D.DiTNoisePredictor.CreateAttentionLayer,UViTNoisePredictorblocks, andVideoUNetPredictor.Create{Spatial,Temporal,Cross}Attentionall switched to the lazy helpers.LayerNormalizationLayerintentionally skipped — gamma/beta are[featureSize]tensors (~18 KB each), not meaningful OOM contributors.Part 3 — Dispose cascades rented weights back to the pool
ConvolutionalLayer.Disposenow callsTensorAllocator.Return(_kernels)in the eager-init branch, matchingDenseLayer.Dispose's pattern for_weights/_biases. Conv kernels are rented viaRentUninitializedin the eager path, so they need to go back for reuse.NeuralNetworkBase.Dispose(bool)cascades to every layer implementingIDisposable. CatchesObjectDisposedExceptionso cross-network layer shares don't abort the loop.DiffusionModelBasegains its ownDispose/Dispose(bool)(it doesn't inherit from NeuralNetworkBase). Subclasses can override to cascade to VAE/text-encoder fields.INeuralNetworkModel<T>andIDiffusionModel<T>now inheritIDisposablesousing var model = CreateModel()compiles. Concrete bases already implemented Dispose — this just tightens the interface surface.DiffusionModelTestBaseandNeuralNetworkModelTestBasewrap all 29CreateModel()/CreateNetwork()call sites inusing var. Combined with PR perf: lazy init in diffusion noise predictors + GC between tests #1137'sIAsyncLifetime.DisposeAsyncGC hook, each test now actively releases weight buffers back to the allocator pool before the next test constructs.Stacking
This PR is branched off
perf/diffusion-lazy-init-oom(#1137's head). When #1137 merges first, the Parts 2+3 diff against master is clean. When this PR merges first and #1137 retains the open state, rebasing is mechanical. Both can land in either order.Build
dotnet build src/AiDotNet.csproj --framework net10.0 -c Release: 0 errorsdotnet build src/AiDotNet.csproj --framework net471 -c Release: 0 errorsdotnet build tests/AiDotNet.Tests/AiDotNetTests.csproj --framework net10.0 -c Release: 0 errorsExpected CI outcome
With all of PR #1137's lazy DenseLayer work AND this PR's lazy attention work AND the Dispose cascade releasing pool tensors between tests, per-instance memory for DiT-XL / SDXL / ControlNet-family models drops from ~4 GB to near-zero at construction, with real allocation scaling to what individual tests actually touch. The three Diffusion ModelFamily shards and the NeuralNetworks shard should transition CANCELLED → SUCCESS or FAILURE.
Refs #1136
Summary by CodeRabbit
Refactor
Chores
Tests