feat: wire JIT compilation through AiModelBuilder

[ooples]

Summary

Plugs the AiDotNet.Tensors compiled-plan infrastructure into AiModelBuilder so a single .ConfigureJitCompilation() call on the builder delivers auto-compiled forward-pass replay to every neural network. The Tensors package has shipped CompiledModelCache, CompiledInferencePlan, CompiledTapeTrainingStep, and the TensorCodec optimization pipeline (dataflow fusion / attention fusion / algebraic backward / CSE / BLAS batching) — but AiDotNet had no builder-level surface to configure it. The builder's YAML schema already documented a jitCompilation section but nothing populated the knob.

Three commits

1. NeuralNetworkBase.Predict virtual default routes through PredictCompiled

Predict changes from abstract to virtual with a default that wraps PredictCompiled(input) in NoGradScope<T>. All 246 existing subclass overrides continue unchanged (their override still binds), so this is a non-breaking API evolution. Subclasses with trivial Predict(x) => Forward(x) bodies can drop their override to pick up compiled inference.

Cache invalidation plumbed into InvalidateParameterCountCache and RestoreOriginalParameters so a lazy-init resize or structure change drops any plan captured against the old layer graph.

2. AiModelBuilder.ConfigureJitCompilation + JitCompilationConfig + result persistence

• New src/Configuration/JitCompilationConfig.cs — POCO mirroring the TensorCodecOptions phase flags (EnableCompilation, EnableDataflowFusion, EnableAttentionFusion, EnableConstantFolding, EnableForwardCSE, EnableBlasBatch, etc.) plus a ThrowOnFailure strict-mode switch. Presets: Default, Aggressive, Disabled. ApplyToTensorCodec() projects onto and installs the Tensors-package options on the current thread.

• AiModelBuilder.ConfigureJitCompilation(JitCompilationConfig?) — follows the ConfigureMixedPrecision / ConfigureReasoning pattern. Config applied at the top of BuildAsync so every subsequent step (training, quantization calibration, cross-validation, RL rollouts) runs under the configured codec flags.

• Config persisted onto AiModelResultOptions and AiModelResult, and re-applied on every Predict call. TensorCodecOptions.Current is [ThreadStatic] — without this bridge, cross-thread inference (request pools, ASP.NET handlers) would see library defaults instead of the builder's flags.

• YAML wiring: added JitCompilation property on YamlModelConfig and YamlConfigApplier binding. The source generator's existingSections set already listed "JitCompilation" so the hand-written property takes precedence cleanly.

• IAiModelBuilder interface updated so every implementation gets the new method signature.

3. Populate JitCompiledFunction so JIT reaches Predict-override models

The previous two commits only set TensorCodecOptions flags. That's enough when a model uses the base NeuralNetworkBase.Predict default (from commit #1). It's NOT enough when a concrete model overrides Predict — and most do: VGG, CNN, FeedForwardNeuralNetwork, etc. return Forward(input) directly, bypassing PredictCompiled. Only AutoTracer (Layer 2, capped at 128 recorded ops) benefits them — too short for deep transformers.

This commit populates the JitCompiledFunction slot on AiModelResult at build time via a new BuildCompiledPredictFunction helper. The helper wraps the model's Predict in a CompiledModelCache regardless of how the concrete override is written. AiModelResult.Predict already checks JitCompiledFunction at line 1877 and routes to the compiled path when set.

Key implementation details:

• Gated on _jitCompilationConfig.Enabled and model is NeuralNetworkBase<T>. Diffusion models (DiffusionModelBase does not extend NeuralNetworkBase) and non-neural models (regression, trees) return null — JIT doesn't engage and the result's Predict stays on its existing path. Diffusion coverage is tracked as a structural refactor in a follow-up PR.
• One CompiledModelCache per wrapper. GetOrCompileInference(Tensor, Action) traces on miss, rebinds input data on hit — no per-call allocation.
• Nested-GraphMode defense: during tracing, temporarily set TensorCodecOptions.EnableCompilation = false so any call into NeuralNetworkBase.PredictCompiled from within the model's own Predict falls through to PredictEager instead of opening a second, conflicting GraphMode scope. Without this guard the inner compile drops the outer trace's ops and replay produces wrong results.
• Applies JIT config to the current thread on every call.
• Try/catch with ThrowOnFailure behavior: by default, compilation failures fall back silently to eager. Setting JitCompilationConfig.ThrowOnFailure = true propagates them — use in tests.

Known limitations (documented inline)

• Diffusion models bypass: DiffusionModelBase doesn't extend NeuralNetworkBase, so the wrapper returns null for diffusion. Closing this requires a structural change (CompiledModelHost<T> composition component) that's scoped as its own follow-up PR.
• Silent wrong-results risk: If a covered model has non-Engine tensor access (direct AsSpan writes, scalar control flow) in its forward, trace succeeds but replay reads stale data. Try/catch catches hard failures; silent divergence won't trigger. Users opting into JIT should flip ThrowOnFailure = true in tests and diff compiled vs eager outputs.

Build verification

• dotnet build src/AiDotNet.csproj --framework net10.0 -c Release: 0 errors
• dotnet build src/AiDotNet.csproj --framework net471 -c Release: 0 errors

Usage

// Enable with defaults (silent fallback to eager on failure)
var result = await new AiModelBuilder<float, Tensor<float>, Tensor<float>>()
    .ConfigureModel(myModel)
    .ConfigureJitCompilation()
    .BuildAsync();

// Aggressive (all fusion/CSE passes on)
await builder.ConfigureJitCompilation(JitCompilationConfig.Aggressive).BuildAsync();

// Strict mode for tests
var cfg = JitCompilationConfig.Default;
cfg.ThrowOnFailure = true;
await builder.ConfigureJitCompilation(cfg).BuildAsync();

Summary by CodeRabbit

• New Features

  • Configurable JIT compilation with a fluent builder option and three presets (Default, Aggressive, Disabled). Settings persist with built models, survive cloning/serialization, and can be loaded from YAML.

• Bug Fixes / Improvements

  • Safer JIT: per-shape failure short-circuiting, automatic fallbacks to eager execution, configurable failure behavior, and invalidation of compiled caches when model structure changes. JIT settings are reapplied per prediction for consistent behavior across threads.

[vercel]

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

2 Skipped Deployments

Project	Deployment	Actions	Updated (UTC)
aidotnet_website	Ignored	Preview	Apr 16, 2026 1:14pm
aidotnet-playground-api	Ignored	Preview	Apr 16, 2026 1:14pm

[coderabbitai]

Note

Reviews paused

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

Use the following commands to manage reviews:

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

Use the checkboxes below for quick actions:

• ▶️ Resume reviews
• 🔍 Trigger review

Walkthrough

Adds builder-configurable JIT compilation: new JitCompilationConfig, builder API to persist/apply it, compile-on-first-predict via CompiledModelCache with tracing/guarding, per-result propagation and thread-local TensorCodecOptions application, and default NeuralNetworkBase.Predict favoring compiled inference with eager fallback.

Changes

Cohort / File(s)	Summary
JIT Compilation Config src/Configuration/JitCompilationConfig.cs	New sealed JitCompilationConfig with presets (Default, Aggressive, Disabled), many optimization/fusion/spectral/mixed-precision flags, Validate() and ApplyToTensorCodec() which projects settings into thread-local TensorCodecOptions.
Builder Integration src/AiModelBuilder.cs	Adds _jitCompilationConfig, fluent ConfigureJitCompilation(...), applies config during BuildAsync, and adds BuildCompiledPredictFunction(...) to produce a JIT-compiled predict delegate (uses CompiledModelCache, nested-compilation guard, and eager fallback on failure).
Builder API & Options src/Interfaces/IAiModelBuilder.cs, src/Models/Options/AiModelResultOptions.cs	Interface gains ConfigureJitCompilation(...); AiModelResultOptions gains JitCompilationConfig?. Builder populates JitCompilationConfig and sets JitCompiledFunction in result options.
Model Result Runtime src/Models/Results/AiModelResult.cs	Stores JitCompilationConfig in result; constructor/deserialization/clones propagate it; Predict(...) reapplies JitCompilationConfig.ApplyToTensorCodec() per call or resets to default when absent.
Neural Network Execution src/NeuralNetworks/NeuralNetworkBase.cs	Predict changed from abstract→virtual with default NoGradScope→PredictCompiled; adds _knownBadCompileShapes, ComputeShapeKey, PredictCompiled that short-circuits known-bad shapes, records failures and logs warnings, and invalidates compiled caches on parameter/structure changes.
YAML Integration src/Configuration/YamlModelConfig.cs, src/Configuration/YamlConfigApplier.cs	YAML model config adds JitCompilation property and applier wires it to builder via ConfigureJitCompilation(...) during configuration application.

Sequence Diagram(s)

sequenceDiagram
    participant Builder as AiModelBuilder
    participant Config as JitCompilationConfig
    participant TCO as TensorCodecOptions
    participant Cache as CompiledModelCache
    participant NN as NeuralNetworkBase
    participant Result as AiModelResult.Predict

    Builder->>Config: ConfigureJitCompilation(config)
    Builder->>Config: config.Validate()
    Builder->>TCO: config.ApplyToTensorCodec() during BuildAsync

    Result->>Config: If present, ApplyToTensorCodec() (thread-local)
    Result->>NN: NN.Predict(input)
    NN->>NN: Enter NoGradScope
    NN->>Cache: Check compiled plan for model/signature

    alt Cache Hit
        Cache->>NN: Return compiled plan
        NN->>NN: Execute compiled plan
    else Cache Miss
        NN->>NN: Trace/record forward pass (with nested-compilation guard)
        NN->>Cache: Compile & store plan (may throw)
        Cache->>NN: Execute compiled plan
    end

    alt Compilation error & ThrowOnFailure == false
        NN->>NN: TraceWarning logged
        NN->>NN: Fallback to PredictEager(input)
    end

Loading

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~50 minutes

Possibly related issues

• perf: JIT compiler integration for zero-allocation diffusion model forward passes #1015: Implements JIT integration, thread-local TensorCodecOptions, compiled-predict wrapper and cache requested by this issue.
• feat: plan serialization — SaveAsync/LoadAsync on ICompiledPlan + ICompiledTrainingPlan AiDotNet.Tensors#166: Adds compiled-plan lifecycle and TensorCodecOptions usage that align with this PR’s compiled-inference caching and codec serialization concerns.

Possibly related PRs

• feat: auto-compiled training integration (Phase 2) #1107: Strong overlap—modifies compiled-inference/JIT pipeline and NeuralNetworkBase predict/cache behavior.
• perf: JIT compiler integration for zero-allocation diffusion forward passes #1016: Related—both extend JIT/codegen/fused execution flows used by compiled delegates.
• refactor: remove per-layer JIT compiler system #1099: Potential conflict—addresses removal of per-layer JIT and JitCompilationConfig which this PR adds.

⚠️ Code Quality Flags (BLOCKING)

1. TODOs/placeholders/stubs: flag every TODO/XXX/placeholder in new and modified files; treat as BLOCKING until implemented or explicitly documented as safe scaffolding.
2. JitCompilationConfig.Validate() exceptions must include descriptive messages (property, actual value, expected range). BLOCKING.
3. Thread-local mutation safety: JitCompilationConfig.ApplyToTensorCodec() sets thread-static TensorCodecOptions without push/pop or restore semantics; concurrent/nested Predicts can corrupt other callers. BLOCKING: implement push/pop, restore, or document strict thread-affinity.
4. Mutable config retained post-Build: builder stores a mutable JitCompilationConfig reference; mutations after Build change runtime behavior unpredictably. BLOCKING: snapshot/clone config when storing in result or make immutable.
5. Silent fallback observability: when compilation fails and ThrowOnFailure == false, only a TraceWarning is emitted; insufficient production diagnostics. BLOCKING: provide structured metrics/logging or require explicit opt-in to silence.
6. API/behavior change: changing NeuralNetworkBase.Predict from abstract to virtual is a behavioral surface change—verify all subclasses still satisfy intended semantics; treat as breaking unless intended. BLOCKING.
7. Exception handling during trace/compile: ensure exceptions during tracing/compile produce clear diagnostics and do not silently swallow important errors; current code may rethrow based on ThrowOnFailure but tracing path needs deterministic handling. BLOCKING.
8. Thread-static global reset path: when JitCompilationConfig is null, code "resets to default" TensorCodecOptions.Current — ensure this is safe in pooled/threaded environments and does not race with other callers. BLOCKING.
9. Serialized config mutability/security: storing JitCompilationConfig in serialized results may capture environment-specific flags (e.g., aggressive optimizations) that are unsafe across platforms; require validation on deserialize. BLOCKING.

Poem

Threads don the flags, tensors hum and weave,
A cache remembers what replays can achieve.
Compile on first sight, fall back if it breaks,
Guards and warnings keep careful wakes.
Snapshot, restore, and logs to fix mistakes. 🎛️

🚥 Pre-merge checks | ✅ 3

✅ Passed checks (3 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title 'feat: wire JIT compilation through AiModelBuilder' clearly and concisely summarizes the primary change: integrating JIT compilation infrastructure into the builder pattern.
Docstring Coverage	✅ Passed	Docstring coverage is 88.00% which is sufficient. The required threshold is 80.00%.

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

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch perf/jit-compilation-wiring

---

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

[Copilot]

Pull request overview

This PR adds a builder-level surface area for configuring and propagating the AiDotNet.Tensors JIT/compiled-plan infrastructure so that a single ConfigureJitCompilation() call can control compiled inference behavior across the built model (including cross-thread inference where TensorCodecOptions.Current is thread-static).

Changes:

• Made NeuralNetworkBase<T>.Predict virtual with a default implementation that routes through compiled inference (PredictCompiled) under NoGradScope<T>, and added compiled-cache invalidation on layer-structure changes.
• Introduced JitCompilationConfig and wired it through AiModelBuilder, IAiModelBuilder, YAML config (jitCompilation), and persistence onto AiModelResultOptions/AiModelResult with per-Predict re-application.
• Added a builder helper intended to populate AiModelResult.JitCompiledFunction via CompiledModelCache to reach models that override Predict.

Reviewed changes

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

Show a summary per file

File	Description
src/NeuralNetworks/NeuralNetworkBase.cs	Default Predict now uses compiled inference under NoGradScope; compiled cache invalidated on structural changes.
src/AiModelBuilder.cs	Adds ConfigureJitCompilation, applies config during BuildAsync, and introduces BuildCompiledPredictFunction wrapper.
src/Configuration/JitCompilationConfig.cs	New config POCO mirroring TensorCodecOptions flags + validation + apply-to-thread behavior.
src/Models/Options/AiModelResultOptions.cs	Persists JitCompilationConfig onto result options.
src/Models/Results/AiModelResult.cs	Stores JIT config and applies it at start of Predict to bridge thread-static options.
src/Interfaces/IAiModelBuilder.cs	Exposes ConfigureJitCompilation on the builder interface.
src/Configuration/YamlModelConfig.cs	Adds JitCompilation section binding for YAML.
src/Configuration/YamlConfigApplier.cs	Applies YAML jitCompilation section to the builder.

---

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

[coderabbitai]

Actionable comments posted: 9

Caution

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

⚠️ Outside diff range comments (1)

src/NeuralNetworks/NeuralNetworkBase.cs (1)

2821-2828: ⚠️ Potential issue | 🟠 Major

Invalidate all structural caches when lazy init reshapes parameters.

This branch only bumps _layerStructureVersion and invalidates _compiledInferenceCache. When lazy initialization changes parameter sizes, _cachedParameterCount and _cachedLayerInfo are stale too, so ParameterCount, GetAllLayerInfo(), and any code that sizes buffers from them can return incorrect metadata after the first training step.

Concrete fix

         if (anyStructureChanged)
         {
-            _parameterBuffer = null;
-            _layerStructureVersion++;
-            // Invalidate the compiled inference cache — the layer graph changed
-            // (e.g., lazy init resized a weight tensor), so any plan traced
-            // before is now pointing at dead tensor references.
-            _compiledInferenceCache?.Invalidate();
+            InvalidateParameterCountCache();
         }

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@src/NeuralNetworks/NeuralNetworkBase.cs` around lines 2821 - 2828, The branch
that runs when anyStructureChanged currently nulls _parameterBuffer, increments
_layerStructureVersion, and invalidates _compiledInferenceCache but does not
clear other derived structural caches; update this block to also
invalidate/reset _cachedParameterCount and _cachedLayerInfo so that
Property/methods like ParameterCount and GetAllLayerInfo recompute fresh
metadata after lazy init reshapes parameters (refer to the symbols
_parameterBuffer, _layerStructureVersion, _compiledInferenceCache,
_cachedParameterCount, _cachedLayerInfo, ParameterCount, and GetAllLayerInfo).

🤖 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/AiModelBuilder.cs`:
- Line 1436: The JIT compiled wrapper is only set on the main supervised result
path causing inference-only, streaming, and RL results to miss
JitCompiledFunction; factor BuildCompiledPredictFunction(...) into the shared
result construction and call it wherever a concrete model result is created
(including the code paths that persist JitCompilationConfig such as the main
builder, inference-only/streaming, and RL result factories referenced around
JitCompilationConfig assignments), then assign the returned compiled delegate to
the result.JitCompiledFunction before returning; update
ConfigureJitCompilation() usages to rely on this shared builder so all
entrypoints (the supervised Predict path and the inference/streaming/RL paths at
the noted locations) receive the same wrapper.
- Around line 1267-1273: The catch block guarded by "catch when
(!throwOnFailure)" that returns new[] { nnModel.Predict(input) } must log the
swallowed exception before falling back; modify this catch to emit a
warning/trace including the exception details, the model type (nnModel.GetType()
or equivalent), and the input shape (e.g., input.Length or input.Shape property)
so JIT/replay failures are visible in logs, then proceed to call
nnModel.Predict(input) as the fallback.
- Around line 1218-1221: The call jitConfig.ApplyToTensorCodec() mutates
thread-static TensorCodecOptions which doesn't flow across async/await and can
leak to other threads; change the state handling so it flows with async control
flow by converting TensorCodecOptions to use AsyncLocal<T> (preferred) or wrap
ApplyToTensorCodec with a save/restore scope around synchronous work so previous
state is restored after use; update BuildAsync and every other call site of
jitConfig.ApplyToTensorCodec() (including the sites around the noted later
calls) to use the new AsyncLocal-backed options or the save/restore pattern so
the JIT config follows the async context and cannot pollute unrelated pooled
threads.

In `@src/Configuration/JitCompilationConfig.cs`:
- Around line 186-205: The ApplyToTensorCodec method sets a new
TensorCodecOptions instance into the thread-local store without validating the
incoming values; call the options' Validate() method on the created
AiDotNet.Tensors.Engines.Optimization.TensorCodecOptions instance (opts) and
only call TensorCodecOptions.SetCurrent(opts) if validation succeeds (propagate
or surface the validation error otherwise) so invalid values like
DataflowFusionMaxHidden <= 0 can't be installed.

In `@src/Interfaces/IAiModelBuilder.cs`:
- Around line 1350-1363: The XML docs for the JIT enablement method overstate
behavior by saying it applies to "forward and backward passes" while the
implementation only auto-compiles and replays the forward (inference) pass;
update the documentation on the method that references
AiDotNet.Configuration.JitCompilationConfig, TensorCodecOptions, Build(),
AiModelResult, and Predict to explicitly state JIT/inlining applies to the
forward/inference pass (first-call trace + replay) and remove or narrow any
claim about backward/pass or training-time compilation so the public API docs
match the wired behavior.

In `@src/Models/Results/AiModelResult.cs`:
- Around line 1280-1281: WithParameters(...) and DeepCopy() construct new
AiModelResultOptions<T, TInput, TOutput> but do not copy the
JitCompilationConfig, causing per-thread JIT settings to be lost; update both
methods so the new AiModelResultOptions instance sets JitCompilationConfig =
this.JitCompilationConfig (or JitCompilationConfig = JitCompilationConfig) when
constructing the copy to propagate the JIT config through the copy/clone flows.
- Around line 644-651: The Deserialize method fails to restore the serialized
JitCompilationConfig onto the current instance, so update Deserialize to assign
the deserialized value back to this instance: inside Deserialize(...) set
this.JitCompilationConfig = deserializedObject.JitCompilationConfig (referencing
the private property JitCompilationConfig and the local deserializedObject) so
reloaded AiModelResult instances reapply the builder JIT flags.
- Around line 1825-1829: When JitCompilationConfig is null the current code
leaves TensorCodecOptions.Current as-is causing stale thread-static codec state
to leak into Predict(); modify the surrounding code in AiModelResult (around the
JitCompilationConfig?.ApplyToTensorCodec() call) to snapshot
TensorCodecOptions.Current before the Predict() work and always restore that
snapshot on exit (use a try/finally) so that if JitCompilationConfig is null
nothing from a previous request remains; if JitCompilationConfig is non-null
still call JitCompilationConfig.ApplyToTensorCodec() and restore the prior
snapshot afterwards to maintain per-thread isolation.

In `@src/NeuralNetworks/NeuralNetworkBase.cs`:
- Around line 2094-2098: The current Predict method routes all calls to
PredictCompiled which unconditionally swallows JIT/compilation failures; update
the implementation so the base/default inference path respects the
ThrowOnFailure contract: change or overload PredictCompiled to accept a boolean
throwOnFailure parameter and call PredictCompiled(input, throwOnFailure: true)
from Predict (still inside NoGradScope<T>), and implement the try/catch inside
PredictCompiled to rethrow when throwOnFailure is true and only log+fallback to
PredictEager(input) when false; reference Predict, PredictCompiled, PredictEager
and the ThrowOnFailure behavior when making the change.

---

Outside diff comments:
In `@src/NeuralNetworks/NeuralNetworkBase.cs`:
- Around line 2821-2828: The branch that runs when anyStructureChanged currently
nulls _parameterBuffer, increments _layerStructureVersion, and invalidates
_compiledInferenceCache but does not clear other derived structural caches;
update this block to also invalidate/reset _cachedParameterCount and
_cachedLayerInfo so that Property/methods like ParameterCount and
GetAllLayerInfo recompute fresh metadata after lazy init reshapes parameters
(refer to the symbols _parameterBuffer, _layerStructureVersion,
_compiledInferenceCache, _cachedParameterCount, _cachedLayerInfo,
ParameterCount, and GetAllLayerInfo).

🪄 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: b225c9f3-096b-41b2-9aba-d41145d4a061

📥 Commits

Reviewing files that changed from the base of the PR and between 72f425a and 07bb43a.

📒 Files selected for processing (8)

• src/AiModelBuilder.cs
• src/Configuration/JitCompilationConfig.cs
• src/Configuration/YamlConfigApplier.cs
• src/Configuration/YamlModelConfig.cs
• src/Interfaces/IAiModelBuilder.cs
• src/Models/Options/AiModelResultOptions.cs
• src/Models/Results/AiModelResult.cs
• src/NeuralNetworks/NeuralNetworkBase.cs

[coderabbitai]

Actionable comments posted: 5

Caution

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

⚠️ Outside diff range comments (1)

src/NeuralNetworks/NeuralNetworkBase.cs (1)

2828-2835: ⚠️ Potential issue | 🟠 Major

Invalidate the metadata caches too when lazy init changes parameter shapes.

When anyStructureChanged is true, this block only drops _parameterBuffer and the compiled plan cache. _cachedParameterCount and _cachedLayerInfo can still describe the pre-init parameter layout, so ParameterCount, SetParameters(), and GetAllLayerInfo() may be wrong after the first training pass.

Proposed fix

         _savedOriginalParameters = null;
 
         // Only invalidate the parameter buffer when layer structure actually changed
         // (e.g., lazy initialization resized a layer). When structure is stable
         // (normal training iterations), keep the buffer to avoid O(total_params)
         // rebuild cost every iteration — critical for large models like VideoCLIP.
         if (anyStructureChanged)
         {
-            _parameterBuffer = null;
-            _layerStructureVersion++;
-            // Invalidate the compiled inference cache — the layer graph changed
-            // (e.g., lazy init resized a weight tensor), so any plan traced
-            // before is now pointing at dead tensor references.
-            _compiledInferenceCache?.Invalidate();
+            InvalidateParameterCountCache();
         }

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@src/NeuralNetworks/NeuralNetworkBase.cs` around lines 2828 - 2835, When
anyStructureChanged is true, also clear the metadata caches so cached
counts/infos don't reflect pre-init shapes: in the same block that sets
_parameterBuffer = null, increments _layerStructureVersion, and calls
_compiledInferenceCache?.Invalidate(), also reset _cachedParameterCount and
_cachedLayerInfo (and any related metadata cache fields) so ParameterCount,
SetParameters(), and GetAllLayerInfo() will recompute based on the new shapes
rather than returning stale data.

♻️ Duplicate comments (2)

src/NeuralNetworks/NeuralNetworkBase.cs (1)

2094-2098: ⚠️ Potential issue | 🔴 Critical

Blocking: the base Predict() path still ignores ThrowOnFailure.

Predict() now funnels every default inference call through PredictCompiled(), but PredictCompiled() always logs and falls back to eager. For models that keep the base implementation, callers cannot opt into fail-fast JIT behavior, so the new config contract is silently bypassed.

Proposed fix

 public virtual Tensor<T> Predict(Tensor<T> input)
 {
     using var _ = new NoGradScope<T>();
-    return PredictCompiled(input);
+    return PredictCompiled(input, throwOnFailure: ShouldThrowOnCompiledPredictFailure());
 }
 
-protected Tensor<T> PredictCompiled(Tensor<T> input)
+protected virtual bool ShouldThrowOnCompiledPredictFailure() => false;
+
+protected Tensor<T> PredictCompiled(Tensor<T> input, bool throwOnFailure)
 {
     if (!AiDotNet.Tensors.Engines.Optimization.TensorCodecOptions.Current.EnableCompilation)
         return PredictEager(input);
 
     try
@@
     }
     catch (Exception ex)
     {
+        if (throwOnFailure)
+            throw;
+
         System.Diagnostics.Trace.TraceWarning(
             $"PredictCompiled fallback for {GetType().FullName} " +
             $"with input shape [{string.Join(", ", input.Shape)}]: {ex.GetType().Name}: {ex.Message}");
         return PredictEager(input);
     }
 }

Also applies to: 2142-2151

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@src/NeuralNetworks/NeuralNetworkBase.cs` around lines 2094 - 2098, The base
Predict method currently always calls PredictCompiled which silently falls back
to eager execution; update Predict (and the analogous overloads at lines
~2142-2151) to honor the ThrowOnFailure setting by checking the model/config
ThrowOnFailure flag inside Predict (while still using NoGradScope<T>), and if
ThrowOnFailure is true call the compiled/JIT path in a way that surfaces
compilation failures (e.g., call a new PredictCompiled overload or a PredictJit
helper that throws on compilation errors) instead of allowing PredictCompiled's
fallback; ensure references to ThrowOnFailure, Predict, PredictCompiled, and
NoGradScope<T> are used so models that don't override the base method will get
fail-fast behavior.

src/AiModelBuilder.cs (1)

1218-1221: ⚠️ Potential issue | 🔴 Critical

Scope or reset TensorCodecOptions instead of leaving thread-static JIT state behind.

Both call sites mutate thread-static codec state and never restore a clean baseline. That creates two correctness hazards: pooled threads can carry JIT flags into unrelated work, and the no-config BuildAsync() path can inherit stale options from a previous request because it never reinstalls defaults. This needs a real scope/reset strategy (or an AsyncLocal-backed implementation in Tensors), not one-way mutation.

Also applies to: 1295-1309

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@src/AiModelBuilder.cs` around lines 1218 - 1221, The thread-static mutation
done by jitConfig.ApplyToTensorCodec() leaves TensorCodecOptions altered across
requests; change this to a scoped/reset pattern so the previous thread-static
state is restored after each use (or replace thread-static with
AsyncLocal-backed storage in the Tensors layer). Specifically, wrap calls that
currently invoke jitConfig.ApplyToTensorCodec() (including the call in
AiModelBuilder.BuildAsync and the other site around lines 1295-1309) in a scope
that captures the current TensorCodecOptions, applies jitConfig, then guarantees
restoration of the original options (e.g., try/finally or IDisposable-style
guard) so pooled threads and no-config BuildAsync() do not inherit stale flags;
alternatively implement TensorCodecOptions accessors to push/pop or use
AsyncLocal to avoid cross-request leakage. Ensure the new scope/reset is used at
both call sites and add tests exercising BuildAsync() with and without jitConfig
to verify no state bleed.

🤖 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/AiModelBuilder.cs`:
- Around line 1198-1201: The code currently wraps every
NeuralNetworks.NeuralNetworkBase<T> with a compiled wrapper, but you must only
install that wrapper for subclasses that override Predict; otherwise calling
nnModel.Predict(input) can re-enter PredictCompiled and break the fallback.
Change the gate around BuildCompiledPredictFunction(...) so it first checks
whether the concrete type overrides the Predict method (i.e., only for
subclasses that provide an override of Predict), and skip building/installing
the compiled wrapper for instances that rely on the base implementation; keep
the existing fallback path in PredictCompiled (the logic around
nnModel.Predict(input) in the block referenced at 1267-1282) unchanged so
base-class networks retain native compiled/eager behavior.
- Around line 1453-1459: BuildMetaLearningInternalAsync currently constructs
AiModelResultOptions<T, TInput, TOutput> without copying JitCompilationConfig
and JitCompiledFunction, which causes ConfigureJitCompilation() to be ignored on
the meta-learning path; fix this by populating those two fields when building
the result (either by moving JitCompilationConfig = _jitCompilationConfig and
JitCompiledFunction =
BuildCompiledPredictFunction(optimizationResult.BestSolution) into the shared
result-construction helper used by BuildAsync, or by explicitly adding those two
assignments in BuildMetaLearningInternalAsync), and audit the other builders
(the spots referenced by the reviewer: the other result-builder locations) to
ensure every AiModelResultOptions construction includes JitCompilationConfig and
JitCompiledFunction.

In `@src/Configuration/JitCompilationConfig.cs`:
- Around line 197-198: The XML doc on JitCompilationConfig mentions
AiModelBuilder{T, TInput, TOutput}.Build(), but the implementation now uses
BuildAsync; update the stale cref and text to reference BuildAsync (e.g., <see
cref="AiModelBuilder{T, TInput, TOutput}.BuildAsync"/>) and any surrounding
wording to reflect the async method; search the same file for other occurrences
of "Build()" referring to AiModelBuilder and replace them with the async variant
to keep API docs accurate.
- Around line 179-183: The validation in JitCompilationConfig for the
SpectralErrorTolerance property currently allows float.NaN to pass; update the
validation block (the code that throws InvalidOperationException for
SpectralErrorTolerance) to explicitly reject NaN by checking
float.IsNaN(SpectralErrorTolerance) (and optionally float.IsInfinity if you want
to be safe) and throw a clear InvalidOperationException that includes the
offending value and states it must be in [0, 1) and not NaN; keep the existing
range checks for <0f and >=1f in the same validation path so all invalid cases
are handled together.

In `@src/Models/Options/AiModelResultOptions.cs`:
- Around line 596-613: Add a missing <value> XML documentation element for the
AiModelResultOptions.JitCompilationConfig property: update the property comment
block in class AiModelResultOptions to include a <value> element that succinctly
describes what the JitCompilationConfig holds/controls (e.g., the JIT
compilation flags applied on each Predict call and how they are used), keeping
the existing <summary> and <remarks> intact to satisfy the documentation
contract requiring <summary>, <value>, and <remarks> for every property.

---

Outside diff comments:
In `@src/NeuralNetworks/NeuralNetworkBase.cs`:
- Around line 2828-2835: When anyStructureChanged is true, also clear the
metadata caches so cached counts/infos don't reflect pre-init shapes: in the
same block that sets _parameterBuffer = null, increments _layerStructureVersion,
and calls _compiledInferenceCache?.Invalidate(), also reset
_cachedParameterCount and _cachedLayerInfo (and any related metadata cache
fields) so ParameterCount, SetParameters(), and GetAllLayerInfo() will recompute
based on the new shapes rather than returning stale data.

---

Duplicate comments:
In `@src/AiModelBuilder.cs`:
- Around line 1218-1221: The thread-static mutation done by
jitConfig.ApplyToTensorCodec() leaves TensorCodecOptions altered across
requests; change this to a scoped/reset pattern so the previous thread-static
state is restored after each use (or replace thread-static with
AsyncLocal-backed storage in the Tensors layer). Specifically, wrap calls that
currently invoke jitConfig.ApplyToTensorCodec() (including the call in
AiModelBuilder.BuildAsync and the other site around lines 1295-1309) in a scope
that captures the current TensorCodecOptions, applies jitConfig, then guarantees
restoration of the original options (e.g., try/finally or IDisposable-style
guard) so pooled threads and no-config BuildAsync() do not inherit stale flags;
alternatively implement TensorCodecOptions accessors to push/pop or use
AsyncLocal to avoid cross-request leakage. Ensure the new scope/reset is used at
both call sites and add tests exercising BuildAsync() with and without jitConfig
to verify no state bleed.

In `@src/NeuralNetworks/NeuralNetworkBase.cs`:
- Around line 2094-2098: The base Predict method currently always calls
PredictCompiled which silently falls back to eager execution; update Predict
(and the analogous overloads at lines ~2142-2151) to honor the ThrowOnFailure
setting by checking the model/config ThrowOnFailure flag inside Predict (while
still using NoGradScope<T>), and if ThrowOnFailure is true call the compiled/JIT
path in a way that surfaces compilation failures (e.g., call a new
PredictCompiled overload or a PredictJit helper that throws on compilation
errors) instead of allowing PredictCompiled's fallback; ensure references to
ThrowOnFailure, Predict, PredictCompiled, and NoGradScope<T> are used so models
that don't override the base method will get fail-fast behavior.

🪄 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: f4359c09-d461-4196-a3f8-67ac7de992c5

📥 Commits

Reviewing files that changed from the base of the PR and between 07bb43a and 6b539a1.

📒 Files selected for processing (6)

• src/AiModelBuilder.cs
• src/Configuration/JitCompilationConfig.cs
• src/Interfaces/IAiModelBuilder.cs
• src/Models/Options/AiModelResultOptions.cs
• src/Models/Results/AiModelResult.cs
• src/NeuralNetworks/NeuralNetworkBase.cs

[Copilot]

Pull request overview

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

---

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

[coderabbitai]

Actionable comments posted: 3

♻️ Duplicate comments (3)

src/AiModelBuilder.cs (3)

1460-1466: ⚠️ Potential issue | 🟠 Major

JIT result wiring is still duplicated, and meta-learning still drops it.

These four constructors now copy the same two assignments, but BuildMetaLearningInternalAsync() still builds AiModelResultOptions<T, TInput, TOutput> without either field. ConfigureJitCompilation() therefore remains inconsistent by build path. Please move this into shared result construction instead of patching each initializer separately.

Also applies to: 1819-1825, 3017-3023, 3512-3518

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@src/AiModelBuilder.cs` around lines 1460 - 1466, The JIT wiring
(JitCompilationConfig and JitCompiledFunction) is duplicated across multiple
constructors and missing in BuildMetaLearningInternalAsync leading to
inconsistent ConfigureJitCompilation behavior; refactor by moving the two
assignments into the shared result construction helper used by all build paths:
modify the factory/constructor path that creates AiModelResultOptions<T, TInput,
TOutput> so it always sets JitCompilationConfig = _jitCompilationConfig and
JitCompiledFunction =
BuildCompiledPredictFunction(optimizationResult.BestSolution) (or equivalent
using the available BestSolution variable), and remove the duplicated
per-constructor assignments in AiModelBuilder where JitCompiledFunction is
currently set; ensure BuildMetaLearningInternalAsync uses the same shared
constructor/helper so meta-learning also receives both fields consistently,
preserving ConfigureJitCompilation semantics.

---

1225-1228: ⚠️ Potential issue | 🔴 Critical

Blocking: restore TensorCodecOptions after applying JIT config.

ApplyToTensorCodec() mutates ambient codec state and never restores it. That leaks JIT settings onto pooled threads, and the BuildAsync write still does not survive the first await. This needs a save/restore scope at each call site or an AsyncLocal<T>-backed implementation in the tensors layer.

#!/bin/bash
# Verify how TensorCodecOptions.Current is implemented and where ambient JIT state is mutated.
rg -n -C2 '(\[ThreadStatic\]|ThreadLocal<|AsyncLocal<|class TensorCodecOptions|TensorCodecOptions\.Current|ApplyToTensorCodec\s*\()' src

Also applies to: 1302-1316

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@src/AiModelBuilder.cs` around lines 1225 - 1228, ApplyToTensorCodec() mutates
ambient TensorCodecOptions (TensorCodecOptions.Current) and never restores it,
leaking JIT config to pooled threads and across awaits; wrap each call site
(e.g., the spots calling jitConfig.ApplyToTensorCodec() and the BuildAsync path)
in a save/restore scope: capture the current TensorCodecOptions.Current before
ApplyToTensorCodec(), call ApplyToTensorCodec(), then ensure in a finally block
you restore the previously captured value back to TensorCodecOptions.Current;
alternatively refactor the tensors layer to use an AsyncLocal-backed
TensorCodecOptions.Current implementation, but for a minimal fix implement the
explicit save/restore around ApplyToTensorCodec() calls (also fix the same
pattern at the other occurrences referenced around the BuildAsync flow).

---

1198-1201: ⚠️ Potential issue | 🟠 Major

Only install this wrapper for models that override Predict.

Base NeuralNetworkBase<T> models already have their own compiled/eager path. Wrapping all of them adds a second cache/compile layer, and the fallback on Line 1289 is not actually eager for base implementations because nnModel.Predict(input) re-enters the base compiled path. Gate this on an actual Predict override.

#!/bin/bash
# Inspect the base Predict definition and the wrapper/fallback call sites.
rg -n -C3 'class NeuralNetworkBase|virtual .* Predict\s*\(|override .* Predict\s*\(|BuildCompiledPredictFunction|nnModel\.Predict\s*\(' src

Also applies to: 1274-1289

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@src/AiModelBuilder.cs` around lines 1198 - 1201, The wrapper is being
installed for every NeuralNetworkBase<T> instance, causing a redundant
compile/cache layer when the derived class doesn't override Predict; change the
gate to only install the wrapper when the concrete model actually overrides
Predict by using reflection: obtain the Predict MethodInfo from
model.GetType().GetMethod("Predict", ...) and check that
method.GetBaseDefinition().DeclaringType != typeof(NeuralNetworkBase<T>) (or
equivalently that the method.DeclaringType != typeof(NeuralNetworkBase<T>));
apply this check where you currently test "model is not NeuralNetworkBase<T>
nnModel" (and the related block around the fallback/BuildCompiledPredictFunction
and nnModel.Predict(input) at the other site) so only models with an overridden
Predict get the wrapper/compiled path.

🤖 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/AiModelBuilder.cs`:
- Around line 1203-1208: The code is mutating a shared model instance by setting
nnModel.ThrowOnJitFailure from AiModelBuilder/_jitCompilationConfig, which leaks
per-build policy into the shared nnModel; instead avoid mutating nnModel: pass
the ThrowOnFailure flag through the call path (e.g., into the PredictCompiled
invocation or into the AiModelResult/model-instance wrapper) or clone a
transient model instance before setting the flag so the shared nnModel is never
modified; update the code paths that rely on ThrowOnJitFailure (PredictCompiled
or the model wrapper used by AiModelResult) to accept an explicit
throwOnJitFailure parameter and use that rather than writing to
nnModel.ThrowOnJitFailure.

In `@src/NeuralNetworks/NeuralNetworkBase.cs`:
- Around line 2077-2098: The public capability surface is inconsistent: Predict
now uses PredictCompiled (JIT/tracing) but SupportsJitCompilation still returns
false and ExportComputationGraph throws; fix by aligning the capability
contract—update SupportsJitCompilation to return true when PredictCompiled is
available and implement ExportComputationGraph to delegate to the
compiled/export routine (or to a new ExportComputationGraphCompiled helper) so
callers accurately detect and use JIT-backed inference; alternatively, if you
prefer no-JIT classes, revert Predict to call PredictEager and update docs to
remove JIT wording. Ensure changes reference Predict, PredictCompiled,
SupportsJitCompilation, and ExportComputationGraph.
- Around line 2151-2163: The catch block in PredictCompiled currently logs and
falls back to PredictEager but doesn't record the failed shape, causing repeated
compile attempts; fix this by computing the same shape/structure key used by the
compiled cache (the per-input shape/version identifier), add that key to the
_jitFailureShapes set when catching the exception (inside the existing catch in
PredictCompiled), and before attempting compilation check _jitFailureShapes and
short-circuit directly to PredictEager(input) if the key is present; also ensure
that any place that calls _compiledInferenceCache?.Invalidate() also clears the
corresponding entries in _jitFailureShapes so retries occur after cache
invalidation.

---

Duplicate comments:
In `@src/AiModelBuilder.cs`:
- Around line 1460-1466: The JIT wiring (JitCompilationConfig and
JitCompiledFunction) is duplicated across multiple constructors and missing in
BuildMetaLearningInternalAsync leading to inconsistent ConfigureJitCompilation
behavior; refactor by moving the two assignments into the shared result
construction helper used by all build paths: modify the factory/constructor path
that creates AiModelResultOptions<T, TInput, TOutput> so it always sets
JitCompilationConfig = _jitCompilationConfig and JitCompiledFunction =
BuildCompiledPredictFunction(optimizationResult.BestSolution) (or equivalent
using the available BestSolution variable), and remove the duplicated
per-constructor assignments in AiModelBuilder where JitCompiledFunction is
currently set; ensure BuildMetaLearningInternalAsync uses the same shared
constructor/helper so meta-learning also receives both fields consistently,
preserving ConfigureJitCompilation semantics.
- Around line 1225-1228: ApplyToTensorCodec() mutates ambient TensorCodecOptions
(TensorCodecOptions.Current) and never restores it, leaking JIT config to pooled
threads and across awaits; wrap each call site (e.g., the spots calling
jitConfig.ApplyToTensorCodec() and the BuildAsync path) in a save/restore scope:
capture the current TensorCodecOptions.Current before ApplyToTensorCodec(), call
ApplyToTensorCodec(), then ensure in a finally block you restore the previously
captured value back to TensorCodecOptions.Current; alternatively refactor the
tensors layer to use an AsyncLocal-backed TensorCodecOptions.Current
implementation, but for a minimal fix implement the explicit save/restore around
ApplyToTensorCodec() calls (also fix the same pattern at the other occurrences
referenced around the BuildAsync flow).
- Around line 1198-1201: The wrapper is being installed for every
NeuralNetworkBase<T> instance, causing a redundant compile/cache layer when the
derived class doesn't override Predict; change the gate to only install the
wrapper when the concrete model actually overrides Predict by using reflection:
obtain the Predict MethodInfo from model.GetType().GetMethod("Predict", ...) and
check that method.GetBaseDefinition().DeclaringType !=
typeof(NeuralNetworkBase<T>) (or equivalently that the method.DeclaringType !=
typeof(NeuralNetworkBase<T>)); apply this check where you currently test "model
is not NeuralNetworkBase<T> nnModel" (and the related block around the
fallback/BuildCompiledPredictFunction and nnModel.Predict(input) at the other
site) so only models with an overridden Predict get the wrapper/compiled path.

🪄 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: e7ff8850-6372-44a7-b2db-04ba697cf5d6

📥 Commits

Reviewing files that changed from the base of the PR and between 6b539a1 and 71b7c8d.

📒 Files selected for processing (2)

• src/AiModelBuilder.cs
• src/NeuralNetworks/NeuralNetworkBase.cs

[coderabbitai]

Actionable comments posted: 1

Caution

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

⚠️ Outside diff range comments (1)

src/AiModelBuilder.cs (1)

1449-1467: 🛠️ Refactor suggestion | 🟠 Major

Factor the JIT result wiring into one helper.

The same JitCompilationConfig/JitCompiledFunction block is now copied into four result builders plus a meta-learning variant. This already regressed once; keeping it duplicated makes the next path-specific drift likely.

♻️ Suggested direction

+    private void ApplyJitResultOptions(
+        AiModelResultOptions<T, TInput, TOutput> options,
+        IFullModel<T, TInput, TOutput>? model)
+    {
+        options.JitCompilationConfig = _jitCompilationConfig;
+        options.JitCompiledFunction = BuildCompiledPredictFunction(model);
+    }

Then call this from each result-construction path, and let the meta-learning path opt out explicitly by passing null if it truly cannot expose a compiled wrapper yet.

Also applies to: 1808-1826, 2999-3024, 3177-3188, 3505-3527

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@src/AiModelBuilder.cs` around lines 1449 - 1467, The JIT wiring
(JitCompilationConfig + JitCompiledFunction = BuildCompiledPredictFunction(...))
is duplicated across multiple result builders; extract this into a single helper
(e.g., CreateJitCompiledWrapper or ConfigureJitForResult) that accepts the
chosen solution/optimizationResult and returns the JitCompiledFunction (or null)
and/or applies JitCompilationConfig, then replace the inline blocks in the
AiModelBuilder result construction sites (references: JitCompilationConfig,
JitCompiledFunction, BuildCompiledPredictFunction, ConfigureJitCompilation) to
call that helper; ensure the meta-learning path explicitly opts out by passing a
flag or null to the helper when a compiled wrapper must not be exposed.

♻️ Duplicate comments (4)

src/NeuralNetworks/NeuralNetworkBase.cs (1)

2097-2100: ⚠️ Potential issue | 🟠 Major

Align the public JIT capability contract with the new default Predict() path.

Predict() now enters PredictCompiled(), but this class still advertises SupportsJitCompilation == false and ExportComputationGraph(...) still throws. Callers that gate compiled inference on the public capability surface will now get a false negative and bypass the behavior this change introduced.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@src/NeuralNetworks/NeuralNetworkBase.cs` around lines 2097 - 2100, The public
Predict() now delegates to PredictCompiled(), but the class still reports
SupportsJitCompilation == false and ExportComputationGraph(...) throws, causing
callers to incorrectly skip compiled inference; update the capability contract
so they agree: in NeuralNetworkBase make the SupportsJitCompilation property
return true by default (or override it to reflect actual compiled support) and
implement/forward ExportComputationGraph(...) to the compiled-export path (or
make it virtual and call the new compiled exporter used by PredictCompiled),
ensuring Predict, PredictCompiled, SupportsJitCompilation, and
ExportComputationGraph are consistent.

src/AiModelBuilder.cs (2)

1226-1230: ⚠️ Potential issue | 🔴 Critical

Restore TensorCodecOptions instead of mutating it one-way.

These calls write thread-static state and never restore it. After the first await, the build can resume on a different thread with default codec options, while the original pooled worker keeps the old JIT flags for unrelated work. Scope the options around the synchronous tensor-execution blocks instead of applying them globally and leaving them behind.

In .NET/C#, do [ThreadStatic] values flow across async/await continuations, and is it safe to leave thread-local state mutated on pooled worker threads after the work item completes?

🔧 Suggested direction

-            jitConfig.ApplyToTensorCodec();
-            try
+            var savedOptions = AiDotNet.Tensors.Engines.Optimization.TensorCodecOptions.Current;
+            jitConfig.ApplyToTensorCodec();
+            try
             {
                 var plan = cache.GetOrCompileInference(input, () =>
                 {
                     ...
                 });

                 return new[] { plan.Execute() };
             }
             catch (Exception ex) when (!throwOnFailure)
             {
                 ...
             }
+            finally
+            {
+                AiDotNet.Tensors.Engines.Optimization.TensorCodecOptions.SetCurrent(savedOptions);
+            }

Apply the same save/restore pattern at each synchronous compute entry point instead of the one-time mutation at the top of BuildAsync.

Also applies to: 1303-1317

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@src/AiModelBuilder.cs` around lines 1226 - 1230, The code calls
jitConfig.ApplyToTensorCodec() which mutates thread-static TensorCodecOptions
and never restores it; change each usage in BuildAsync (and the similar spots
around 1303-1317) to a save/set/restore pattern: capture the current
TensorCodecOptions, apply jitConfig for the immediate synchronous tensor
execution, execute the synchronous work, and restore the previously saved
TensorCodecOptions in a finally block so pooled worker threads don't keep the
JIT flags; wrap only the synchronous compute entry points (not the whole async
method) with this try/finally restore logic and reuse the same pattern wherever
ApplyToTensorCodec() is invoked.

---

1198-1201: ⚠️ Potential issue | 🟠 Major

Only build this wrapper for concrete Predict overrides.

The fallback on Lines 1289-1290 calls nnModel.Predict(input). For models that still inherit NeuralNetworkBase<T>.Predict, that just routes back into PredictCompiled, so a wrapper failure can immediately retry compiled inference instead of taking a clean eager path. Gate BuildCompiledPredictFunction(...) to types that actually override Predict.

In C#, if a derived type does not override a virtual method, does invoking that method on the derived instance execute the base implementation?

🔧 Suggested fix

         if (model is not NeuralNetworks.NeuralNetworkBase<T> nnModel) return null;
+        var predictMethod = nnModel.GetType().GetMethod(
+            nameof(NeuralNetworks.NeuralNetworkBase<T>.Predict),
+            new[] { typeof(Tensor<T>) });
+        if (predictMethod?.DeclaringType == typeof(NeuralNetworks.NeuralNetworkBase<T>))
+        {
+            return null;
+        }
 
         // ThrowOnFailure is enforced at THIS wrapper level (the catch below).

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@src/AiModelBuilder.cs` around lines 1198 - 1201, The current logic builds a
compiled wrapper even when the concrete model type does not override Predict,
causing the fallback to call nnModel.Predict(input) which can re-enter
PredictCompiled; change BuildCompiledPredictFunction so it only proceeds for
concrete types that actually override Predict: detect whether the runtime type
of nnModel has an override of the virtual Predict method (compare the runtime
type's MethodInfo for Predict against NeuralNetworkBase<T>.Predict and
skip/return null if DeclaringType equals NeuralNetworkBase<T>), and leave
PredictCompiled/Predict behavior unchanged for types that don't override.

src/Configuration/JitCompilationConfig.cs (1)

203-204: ⚠️ Potential issue | 🟡 Minor

Update stale XML docs to BuildAsync.

Line 203 still references AiModelBuilder{T, TInput, TOutput}.Build(), but this config is applied from the async build path. Please update the cref/text to BuildAsync to avoid misleading API docs.

📝 Proposed fix

-    /// Called by <see cref="AiModelBuilder{T, TInput, TOutput}"/>.Build() so that
+    /// Called by <see cref="AiModelBuilder{T, TInput, TOutput}.BuildAsync"/> so that

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@src/Configuration/JitCompilationConfig.cs` around lines 203 - 204, Update the
stale XML documentation on JitCompilationConfig so the cref and descriptive text
reference the async build method instead of the sync one: change the <see
cref="AiModelBuilder{T, TInput, TOutput}.Build()"/> reference to <see
cref="AiModelBuilder{T, TInput, TOutput}.BuildAsync()"/> and update any
surrounding text that mentions "Build()" to "BuildAsync" so the docs accurately
reflect that JitCompilationConfig is applied from the async build path.

🤖 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/NeuralNetworkBase.cs`:
- Around line 2875-2879: When anyStructureChanged is true, do not only call
_compiledInferenceCache?.Invalidate(); instead run the full structural
invalidation sequence used elsewhere: clear or reset _knownBadCompileShapes,
reset _cachedParameterCount to its invalid sentinel, and then invalidate
_compiledInferenceCache so all structural caches reflect lazy-resized
parameters; update the branch that checks anyStructureChanged to perform those
three steps (clear _knownBadCompileShapes, reset _cachedParameterCount,
_compiledInferenceCache?.Invalidate()) rather than only invalidating the
compiled inference cache.

---

Outside diff comments:
In `@src/AiModelBuilder.cs`:
- Around line 1449-1467: The JIT wiring (JitCompilationConfig +
JitCompiledFunction = BuildCompiledPredictFunction(...)) is duplicated across
multiple result builders; extract this into a single helper (e.g.,
CreateJitCompiledWrapper or ConfigureJitForResult) that accepts the chosen
solution/optimizationResult and returns the JitCompiledFunction (or null) and/or
applies JitCompilationConfig, then replace the inline blocks in the
AiModelBuilder result construction sites (references: JitCompilationConfig,
JitCompiledFunction, BuildCompiledPredictFunction, ConfigureJitCompilation) to
call that helper; ensure the meta-learning path explicitly opts out by passing a
flag or null to the helper when a compiled wrapper must not be exposed.

---

Duplicate comments:
In `@src/AiModelBuilder.cs`:
- Around line 1226-1230: The code calls jitConfig.ApplyToTensorCodec() which
mutates thread-static TensorCodecOptions and never restores it; change each
usage in BuildAsync (and the similar spots around 1303-1317) to a
save/set/restore pattern: capture the current TensorCodecOptions, apply
jitConfig for the immediate synchronous tensor execution, execute the
synchronous work, and restore the previously saved TensorCodecOptions in a
finally block so pooled worker threads don't keep the JIT flags; wrap only the
synchronous compute entry points (not the whole async method) with this
try/finally restore logic and reuse the same pattern wherever
ApplyToTensorCodec() is invoked.
- Around line 1198-1201: The current logic builds a compiled wrapper even when
the concrete model type does not override Predict, causing the fallback to call
nnModel.Predict(input) which can re-enter PredictCompiled; change
BuildCompiledPredictFunction so it only proceeds for concrete types that
actually override Predict: detect whether the runtime type of nnModel has an
override of the virtual Predict method (compare the runtime type's MethodInfo
for Predict against NeuralNetworkBase<T>.Predict and skip/return null if
DeclaringType equals NeuralNetworkBase<T>), and leave PredictCompiled/Predict
behavior unchanged for types that don't override.

In `@src/Configuration/JitCompilationConfig.cs`:
- Around line 203-204: Update the stale XML documentation on
JitCompilationConfig so the cref and descriptive text reference the async build
method instead of the sync one: change the <see cref="AiModelBuilder{T, TInput,
TOutput}.Build()"/> reference to <see cref="AiModelBuilder{T, TInput,
TOutput}.BuildAsync()"/> and update any surrounding text that mentions "Build()"
to "BuildAsync" so the docs accurately reflect that JitCompilationConfig is
applied from the async build path.

In `@src/NeuralNetworks/NeuralNetworkBase.cs`:
- Around line 2097-2100: The public Predict() now delegates to
PredictCompiled(), but the class still reports SupportsJitCompilation == false
and ExportComputationGraph(...) throws, causing callers to incorrectly skip
compiled inference; update the capability contract so they agree: in
NeuralNetworkBase make the SupportsJitCompilation property return true by
default (or override it to reflect actual compiled support) and
implement/forward ExportComputationGraph(...) to the compiled-export path (or
make it virtual and call the new compiled exporter used by PredictCompiled),
ensuring Predict, PredictCompiled, SupportsJitCompilation, and
ExportComputationGraph are consistent.

🪄 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: b9bcb7bf-ac8c-48d5-8310-e21da13b134a

📥 Commits

Reviewing files that changed from the base of the PR and between 71b7c8d and 5e4cd1c.

📒 Files selected for processing (5)

• src/AiModelBuilder.cs
• src/Configuration/JitCompilationConfig.cs
• src/Interfaces/IAiModelBuilder.cs
• src/Models/Options/AiModelResultOptions.cs
• src/NeuralNetworks/NeuralNetworkBase.cs

[coderabbitai]

Actionable comments posted: 2

🤖 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/NeuralNetworkBase.cs`:
- Around line 2127-2134: The HashSet<long> field _knownBadCompileShapes is not
thread-safe and can race when Predict/PredictCompiled are called concurrently;
replace it with a thread-safe structure or protect access: either change
_knownBadCompileShapes to a ConcurrentDictionary<long, byte> (use
TryAdd/TryGetValue) or wrap the check-and-add sequence in a lock used by
PredictCompiled/PredictEager, ensuring all contains/add operations on
_knownBadCompileShapes are done via the chosen thread-safe mechanism so
concurrent compilation failures don't corrupt the collection.
- Around line 2080-2095: The XML docs for Predict claim "compiled inference" via
PredictCompiled but the class still reports SupportsJitCompilation = false and
ExportComputationGraph throws, which confuses two different JIT mechanisms
(tracing via CompiledModelCache<T> vs explicit graph export via IJitCompilable);
either add a new capability property (e.g., SupportsTracingCompilation or
SupportsCompiledReplay) and set/describe it where Predict/PredictCompiled are
implemented, or update the Predict XML comments to explicitly state that
PredictCompiled uses tracing/CompiledModelCache<T> and is separate from
IJitCompilable/ExportComputationGraph; update/mention the symbols Predict,
PredictCompiled, SupportsJitCompilation, ExportComputationGraph, IJitCompilable,
and CompiledModelCache<T> so callers/readers can tell which mechanism is
supported.

🪄 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: 86067559-e37e-4975-93c8-f8ec985443d6

📥 Commits

Reviewing files that changed from the base of the PR and between 5e4cd1c and a48649f.

📒 Files selected for processing (1)

• src/NeuralNetworks/NeuralNetworkBase.cs

[Copilot]

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

[ooples]

Refs #1015 — this PR adds the ConfigureJitCompilation builder surface for the JIT compiler integration tracked in issue #1015. Contributes to the "NeuralNetworkBase.CompileForward()" checklist item by exposing the existing Tensors-side CompiledModelCache/CompiledInferencePlan infrastructure through a fluent builder entry point. Does NOT close #1015 — that issue also covers IR operations (GroupNormOp, FusedGroupNormActivationOp), fusion patterns (Patterns 11-14), optimization passes (memory planning, tile scheduling), and benchmark infrastructure — all tracked separately.