chore: jIT Compilation for Autodiff Computation Graphs

[ooples]

PR Title (Auto-Fixed)

Note: PR titles are automatically fixed to follow Conventional Commits format for automated releases.

The workflow will intelligently detect the appropriate type based on:

• Title keywords (fix/add/implement/update/etc.)
• Files changed (docs/tests/ci/source files)
• Default to chore: if unsure

If the auto-detected type is incorrect, simply edit the PR title manually.

User Story / Context

• Reference: [US-XXX] (if applicable)
• Base branch: merge-dev2-to-master

Summary

• What changed and why (scoped strictly to the user story / PR intent)

Verification

• Builds succeed (scoped to changed projects)
• Unit tests pass locally
• Code coverage >= 90% for touched code
• Codecov upload succeeded (if token configured)
• TFM verification (net46, net6.0, net8.0) passes (if packaging)
• No unresolved Copilot comments on HEAD

Copilot Review Loop (Outcome-Based)

Record counts before/after your last push:

• Comments on HEAD BEFORE: [N]
• Comments on HEAD AFTER (60s): [M]
• Final HEAD SHA: [sha]

Files Modified

• List files changed (must align with scope)

Notes

• Any follow-ups, caveats, or migration details

[coderabbitai]

Warning

Rate limit exceeded

@ooples has exceeded the limit for the number of commits or files that can be reviewed per hour. Please wait 11 minutes and 22 seconds before requesting another review.

⌛ How to resolve this issue?

After the wait time has elapsed, a review can be triggered using the @coderabbitai review command as a PR comment. Alternatively, push new commits to this PR.

We recommend that you space out your commits to avoid hitting the rate limit.

🚦 How do rate limits work?

CodeRabbit enforces hourly rate limits for each developer per organization.

Our paid plans have higher rate limits than the trial, open-source and free plans. In all cases, we re-allow further reviews after a brief timeout.

Please see our FAQ for further information.

📥 Commits

Reviewing files that changed from the base of the PR and between 124dfbe and 6af39ee.

📒 Files selected for processing (3)

• docs/JIT_IMPLEMENTATION_STATUS.md (1 hunks)
• src/Autodiff/TensorOperations.cs (1 hunks)
• src/NeuralNetworks/NeuralNetworkBase.cs (2 hunks)

Note

Other AI code review bot(s) detected

CodeRabbit has detected other AI code review bot(s) in this pull request and will avoid duplicating their findings in the review comments. This may lead to a less comprehensive review.

Summary by CodeRabbit

Release Notes

• New Features

  • Added JIT compilation support for accelerated model inference with automatic graph optimization.
  • Enabled configurable optimization passes including constant folding, dead code elimination, and operation fusion.
  • Added compilation caching to avoid redundant compilations.
  • Integrated JIT support across regression, neural network, and time series models.

• Documentation

  • Added comprehensive JIT compiler usage guides, API reference, and configuration documentation.
  • Added example implementations demonstrating common JIT compilation scenarios.

Pre-merge checks and finishing touches

✅ Passed checks (3 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly summarizes the main change: introducing JIT compilation for autodiff computation graphs, which aligns with the comprehensive changeset adding JIT infrastructure, IR operations, optimizations, and model integrations.
Description check	✅ Passed	The description provides context about JIT compilation and references a standard verification checklist, though it is partially a template. It relates to the changeset's goal of adding JIT compilation capabilities.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.

---

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

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

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

[github-actions]

🤖 PR Title Auto-Fixed

Your PR title was automatically updated to follow Conventional Commits format.

Original title:
JIT Compilation for Autodiff Computation Graphs

New title:
chore: jIT Compilation for Autodiff Computation Graphs

Detected type: chore: (default type)
Version impact: No release

---

Valid types and their effects:

• feat: - New feature (MINOR bump: 0.1.0 → 0.2.0)
• fix: - Bug fix (MINOR bump)
• docs: - Documentation (MINOR bump)
• refactor: - Code refactoring (MINOR bump)
• perf: - Performance improvement (MINOR bump)
• test: - Tests only (no release)
• chore: - Build/tooling (no release)
• ci: - CI/CD changes (no release)
• style: - Code formatting (no release)

If the detected type is incorrect, you can manually edit the PR title.

[Copilot]

Pull Request Overview

This PR introduces a JIT (Just-In-Time) compiler for autodiff computation graphs, providing significant performance improvements (5-10x) for neural network inference. The implementation includes a complete IR infrastructure, optimization passes (constant folding, dead code elimination, operation fusion), code generation, and comprehensive testing/benchmarking.

Key changes:

• Complete JIT compiler infrastructure with IR graph representation and 43+ operation types
• Three optimization passes: constant folding, dead code elimination, and operation fusion
• Integration with PredictionModelBuilder for seamless model compilation
• Comprehensive test suite with 90+ unit tests and performance benchmarks
• Extensive documentation with beginner-friendly explanations

Reviewed Changes

Copilot reviewed 41 out of 41 changed files in this pull request and generated 32 comments.

Show a summary per file

File	Description
tests/AiDotNet.Tests/UnitTests/JitCompiler/OptimizationPassTests.cs	Tests for optimization passes (DCE, fusion, constant folding)
tests/AiDotNet.Tests/UnitTests/JitCompiler/JitCompilerTests.cs	Tests for main JIT compiler functionality and caching
tests/AiDotNet.Tests/UnitTests/JitCompiler/IRBuilderTests.cs	Tests for IR graph construction from computation nodes
tests/AiDotNet.Tests/Benchmarks/JitCompilerBenchmarks.cs	Performance benchmarks comparing JIT vs interpreted execution
tests/AiDotNet.Tests/Benchmarks/JIT_BENCHMARKS_README.md	Comprehensive documentation for running and interpreting benchmarks
src/PredictionModelBuilder.cs	Added JIT compilation configuration and integration
src/Models/Results/PredictionModelResult.cs	Added JIT-compiled function storage and execution path
src/Models/NeuralNetworkModel.cs	Added TODO documentation for future JIT support
src/JitCompiler/README.md	High-level architecture and usage documentation
src/JitCompiler/JitCompiler.cs	Main JIT compiler implementation with caching
src/JitCompiler/IRBuilder.cs	Converts ComputationNode graphs to IR representation
src/JitCompiler/Optimizations/*.cs	Optimization pass implementations (DCE, fusion, constant folding, etc.)
src/JitCompiler/IR/*.cs	IR infrastructure (operations, graphs, types, shapes)
src/JitCompiler/CodeGen/*.cs	Code generation utilities (SIMD, gradients)
src/Autodiff/ComputationNode.cs	Added OperationType and OperationParams for JIT metadata

---

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

[Copilot]

This foreach loop implicitly filters its target sequence - consider filtering the sequence explicitly using '.Where(...)'.

[Copilot]

This foreach loop implicitly filters its target sequence - consider filtering the sequence explicitly using '.Where(...)'.

[Copilot]

This foreach loop implicitly filters its target sequence - consider filtering the sequence explicitly using '.Where(...)'.

[Copilot]

This foreach loop implicitly filters its target sequence - consider filtering the sequence explicitly using '.Where(...)'.

[Copilot]

This foreach loop implicitly filters its target sequence - consider filtering the sequence explicitly using '.Where(...)'.

[Copilot]

Generic catch clause.

Suggested change

	{
	{
	// Rethrow critical exceptions that should not be caught
	if (ex is OutOfMemoryException || ex is StackOverflowException || ex is System.Threading.ThreadAbortException)
	throw;

[Copilot]

Both branches of this 'if' statement write to the same variable - consider using '?' to express intent better.

Suggested change

	if (Vector.IsHardwareAccelerated)
	{
	// Vector<T>.Count gives us the number of elements that fit in a SIMD register
	// This is typically 4 for float (128-bit SSE), 8 for AVX, or 16 for AVX-512
	_vectorSize = Vector<float>.Count;
	}
	else
	{
	_vectorSize = 1; // No SIMD support
	}
	// Vector<T>.Count gives us the number of elements that fit in a SIMD register
	// This is typically 4 for float (128-bit SSE), 8 for AVX, or 16 for AVX-512
	// If no SIMD support, use 1
	_vectorSize = Vector.IsHardwareAccelerated
	? Vector<float>.Count
	: 1;

[Copilot]

Both branches of this 'if' statement return - consider using '?' to express intent better.

Suggested change

	if (inputIndex == 0)
	{
	// Gradient to left input (minuend)
	return gradOutput;
	}
	else
	{
	// Gradient to right input (subtrahend) is negated
	return TensorOperations.Negate(gradOutput);
	}
	// Gradient to left input (minuend) is gradOutput; to right input (subtrahend) is -gradOutput
	return inputIndex == 0
	? gradOutput
	: TensorOperations.Negate(gradOutput);

[Copilot]

Both branches of this 'if' statement write to the same variable - consider using '?' to express intent better.

Suggested change

	if (dim1 == -1 || dim2 == -1)
	{
	resultShape[maxRank - i] = -1; // Dynamic
	}
	else
	{
	resultShape[maxRank - i] = Math.Max(dim1, dim2);
	}
	resultShape[maxRank - i] = (dim1 == -1 || dim2 == -1) ? -1 : Math.Max(dim1, dim2);

[Copilot]

Both branches of this 'if' statement write to the same variable - consider using '?' to express intent better.

[coderabbitai]

Actionable comments posted: 17

🧹 Nitpick comments (33)

src/Autodiff/ComputationNode.cs (1)

163-186: Consider thread-safety and immutability for OperationParams.

The Dictionary<string, object>? property is mutable and not thread-safe. If computation nodes are accessed concurrently or if immutability is desired, consider:

1. Using IReadOnlyDictionary<string, object>? for the property type
2. Initializing as a new dictionary in the constructor
3. Making it immutable after construction

However, if the current mutable design aligns with the broader ComputationNode pattern (which already has mutable properties like Value, Gradient, etc.), this is acceptable.

tests/AiDotNet.Tests/Benchmarks/JitCompilerBenchmarks.cs (1)

15-47: Add null-forgiving operators or null checks in benchmark methods.

The benchmark fields are declared as nullable (e.g., _jit?, _simpleGraph?, _simpleCompiled?) but are dereferenced with ! operators in the benchmark methods without validation. While GlobalSetup ensures initialization, consider either:

1. Making fields non-nullable (since GlobalSetup always runs first)
2. Adding explicit null checks in benchmarks for defensive programming

For benchmarks where setup guarantees initialization, non-nullable fields are preferable:

-    private global::AiDotNet.JitCompiler.JitCompiler? _jit;
+    private global::AiDotNet.JitCompiler.JitCompiler _jit = null!;

src/JitCompiler/CodeGen/SIMDOptimizer.cs (4)

84-110: AddSIMDHints is currently a no-op pass-through.

The method returns the input expression unchanged (line 109). While the documentation explains this is intentional (the .NET JIT handles vectorization), consider:

1. Documenting that this is a future extension point
2. Adding a comment explaining why no transformation is applied
3. Potentially removing this method until actual SIMD hints are implemented

The current implementation is functionally correct but may be confusing to future maintainers.

---

65-82: Hardcoded vectorization threshold may need tuning.

Line 78 uses a hardcoded threshold of _vectorSize * 4 to determine if a tensor is "large enough" for SIMD. This threshold:

• May be too conservative for some operations
• May not account for operation complexity
• Could benefit from being configurable

Consider making this threshold configurable through constructor parameters or making it operation-specific based on empirical performance data.

---

115-128: OpType string comparison is case-sensitive and fragile.

The IsElementWiseOp method uses string equality checks (lines 117-127) which are:

• Case-sensitive (may cause mismatches)
• Fragile to typos
• Not validated against actual IROp types

Consider alternatives:

1. Using type-based checks: op is AddOp or SubtractOp or ...
2. Adding a property/interface on IROp: bool IsElementWise { get; }
3. Using string comparison with StringComparison.OrdinalIgnoreCase

The type-based approach would be more robust and catch errors at compile time.

---

175-186: Document the magic number in speedup calculation.

Line 183 uses 0.75 to account for SIMD overhead in the estimated speedup calculation. This appears to be a heuristic, but it's unclear:

• Where this value comes from (empirical data? conservative estimate?)
• If it varies by operation type
• If it accounts for memory bandwidth limitations

Consider:

1. Adding a constant with a descriptive name: const double SIMDEfficiencyFactor = 0.75;
2. Documenting the rationale in a comment
3. Making it configurable if different operations have different efficiencies

docs/JIT-INTEGRATION-SUMMARY.md (4)

83-106: Add explicit language to non-code fenced blocks (MD040).

The integration-flow diagram fenced block has no language spec; markdownlint flags this (MD040). Consider marking it as text (or similar) so it’s clear it’s illustrative, not executable code.

-```
+```text
 ...
-```
+```

---

387-398: Add language to performance characteristics fenced block (MD040).

Same as above: this block is prose/ASCII, not code. Adding a language (e.g., text) will satisfy markdownlint and clarify intent.

-```
+```text
 ...
-```
+```

---

188-238: Avoid duplicate “Completed ✅” headings (MD024).

There are two ### Completed ✅ headings (lines ~188 and ~230), which triggers MD024 and makes navigation ambiguous. Suggest renaming the second one to clarify scope (e.g., “### Backward Pass Completed ✅”):

-### Completed ✅
+### Backward Pass Completed ✅

---

164-184: Keep JitCompilerOptions snippet in sync with the actual API.

The JitCompilerOptions class snippet only lists the four core flags, while later sections and examples refer to additional flags (EnableLoopUnrolling, EnableSIMDHints, EnableAutoTuning, EnableAdaptiveFusion). Please confirm the actual class shape and either:

• Update the snippet to include all current properties, or
• Clearly separate “core” options from “advanced” ones and document both.

This avoids docs drifting from the code and example snippets that won’t compile.

Also applies to: 301-314

docs/JIT-Compiler-Usage-Guide.md (1)

208-212: Minor wording polish for “Very small operations”

The “Less beneficial for” bullet Very small operations (compilation overhead) can be tightened to avoid the overused intensifier (“very”), e.g., “tiny operations” or simply “small operations”.

Purely editorial; no functional impact.

src/Configuration/JitCompilationConfig.cs (1)

62-101: Clarify “all optimizations enabled” vs experimental flags

The CompilerOptions summary/remarks say “all optimizations enabled”, but JitCompilerOptions also exposes experimental flags (e.g., EnableLoopUnrolling, EnableAdaptiveFusion, EnableAutoTuning, EnableSIMDHints) that default to false.

To avoid confusion, consider rephrasing to something like “all core/supported optimizations enabled by default” so it’s clear the experimental ones remain off unless explicitly enabled.

src/Interfaces/IJitCompilable.cs (1)

1-4: Confirm List<> namespace import for IJitCompilable

ExportComputationGraph uses List<ComputationNode<T>>, but this file only imports AiDotNet.Autodiff. Unless you rely on a global using for System.Collections.Generic, this will not compile.

If there’s no global using, add:

+using System.Collections.Generic;
 using AiDotNet.Autodiff;

If a global using exists, no change is needed.

src/JitCompiler/CodeGen/CodeGenerator.cs (1)

1-7: Verify required using directives (LINQ and collections)

This file uses several types/extensions that are not covered by the current usings:

• Dictionary<,>, List<> → System.Collections.Generic
• FirstOrDefault, ToArray, Select → System.Linq
• NotImplementedException, InvalidOperationException → System

If these aren’t provided via project‑wide global usings, you’ll need:

+using System;
+using System.Collections.Generic;
+using System.Linq;
 using System.Linq.Expressions;
 using System.Reflection;

If global usings already exist for these namespaces, you can ignore this.

examples/JitCompiler/BasicUsageExample.cs (3)

131-186: Align “Performance Comparison” example with its description

PerformanceComparisonExample is titled and documented as “Performance comparison (JIT vs interpreted)” but currently only measures the JIT-compiled path.

To avoid confusing users running this example, consider either:

• Adding a baseline interpreted implementation (e.g., evaluate Exp(ReLU(input)) via the non‑JIT path and report timings alongside JIT), or
• Renaming the example/heading to make it clear that it demonstrates raw JIT performance rather than a direct comparison.

This will make the example’s output match reader expectations.

---

11-12: Optional: Make BasicUsageExample a static class

All members of BasicUsageExample are static, and the class is used purely as a container for sample entry points. Marking the class itself as static communicates this intent and prevents accidental instantiation:

-public class BasicUsageExample
+public static class BasicUsageExample

Purely stylistic, but matches common patterns for example/utility holders.

---

1-5: Confirm System.Collections.Generic import for List<> usage

This file uses List<ComputationNode<float>> in multiple places but only imports AiDotNet.Autodiff, AiDotNet.JitCompiler, System, and System.Diagnostics. Unless you rely on a global using for System.Collections.Generic, this will not compile.

If needed, add:

 using AiDotNet.Autodiff;
 using AiDotNet.JitCompiler;
 using System;
 using System.Diagnostics;
+using System.Collections.Generic;

If global usings already cover this, no change is required.

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

349-372: JIT path is safely integrated but has some API and behavior nuances

The JIT hook in Predict is guarded well: it only runs when the normalized input is a Tensor<T> and when the first JIT output is assignable to TOutput, otherwise it cleanly falls back to Model.Predict. That avoids type explosions or normalization mismatches.

A few things to be aware of:

• Generic constraints of the JIT path: As written, JIT is effectively used only for models where TInput normalizes to a single Tensor<T> and TOutput is (or derives from) Tensor<T>. Vector/Matrix or multi‑output scenarios will silently bypass JIT and use the normal path. If the intent is broader JIT coverage, you may want an adapter layer that converts Tensor<T>[] into the appropriate TOutput shape instead of relying on is TOutput on jitResult[0].
• JIT delegate lifetime on copies: WithParameters and DeepCopy both construct new PredictionModelResult instances without passing JitCompiledFunction, so the copies lose JIT acceleration even though they share the same computation graph. That might be acceptable, but if you expect copies to stay JIT‑enabled, consider threading JitCompiledFunction through those constructor calls (being careful about any assumptions the JIT makes about parameter shapes).
• Public constructor signature change: Adding the optional jitCompiledFunction parameter to the long public constructor changes its CLR signature. Any existing binaries compiled against the previous signature will fail to bind. If you care about binary compatibility, adding a new overload instead of extending the existing one would be safer.

None of these are correctness bugs, but they are worth double‑checking against your intended JIT usage and versioning guarantees.

Also applies to: 429-460, 626-663

src/JitCompiler/Optimizations/DeadCodeEliminationPass.cs (1)

40-164: Dead code elimination logic is correct; consider minor refactors for reuse/perf

The pass correctly performs backward liveness from OutputIds, keeps only ops whose OutputId is live, and preserves input/output IDs and tensor shapes, with useful DCE metadata. IdentifyDeadCode and GetStatistics are consistent with this behavior.

Two small follow‑ups you might consider:

• Extract the liveness computation into a private helper used by both Optimize and IdentifyDeadCode to avoid duplicating the fixed‑point loop.
• Replace the outer fixed‑point loop with a standard worklist (queue/stack of newly‑live tensor IDs) so you traverse the ops list only once; current implementation is fine for modest graphs but could become O(N²) on very deep graphs.

These are non‑blocking polish items; functionally the pass looks good.

Also applies to: 166-257

src/JitCompiler/IR/Operations/ActivationOps.cs (1)

3-155: Activation IR ops are consistent with existing IROp patterns

Each activation op correctly enforces a single input via Validate, and ApplyActivationOp additionally guards against an empty ActivationName. The ToString implementations for SoftmaxOp and ApplyActivationOp match the style used in other ops and should be helpful for debugging/IR dumps.

If you later need stricter validation, you could also sanity‑check SoftmaxOp.Axis against OutputShape (or an inferred rank), but that’s an optional enhancement.

tests/AiDotNet.Tests/UnitTests/JitCompiler/JitCompilerTests.cs (1)

89-114: Cache‑hit compilation time assertion may be too strict

Compile_SecondTime_HitsCacheOptimized requires stats2.CompilationTime == TimeSpan.Zero. That’s fine if JitCompiler intentionally sets compilation time to exactly zero on cache hits; however, if you ever change stats collection to include even the cache lookup time, this equality will become brittle.

Consider instead asserting the semantic intent, e.g.:

• Assert.True(stats2.CacheHit); (already present) and
• Assert.True(stats2.CompilationTime <= TimeSpan.Zero || stats2.CompilationTime.TotalMilliseconds < 1);

to allow trivial non‑zero timings while still guaranteeing “effectively free” cached compilations.

src/JitCompiler/IRBuilder.cs (1)

378-423: Topological sort duplication and potential recursion depth

TopologicalSort<T> reimplements a recursive DFS over ComputationNode<T> even though ComputationNode<T> already exposes its own TopologicalSort() helper (iterative stack-based) in the autodiff layer.

While this works functionally, it:

• Duplicates traversal logic already tested on ComputationNode<T>.
• Uses recursion, which is more fragile for very deep graphs than the existing iterative implementation.

Consider delegating to the node’s existing topological sort API, e.g.:

private List<ComputationNode<T>> TopologicalSort<T>(ComputationNode<T> outputNode)
{
    return outputNode.TopologicalSort();
}

That would remove duplication and avoid recursion limits for large graphs.

Also applies to: 84-112

src/JitCompiler/Optimizations/LoopUnrollingPass.cs (1)

84-141: LoopUnrollingPass is currently a no-op scaffold; consider clarifying semantics

As implemented:

• FindRepeatingPattern/AreSimilarOperations/ShouldUnroll detect simple repeated element-wise patterns.
• UnrollPattern just returns a copy of the same operations, without changing IDs or structure.
• Optimize then builds a new IRGraph whose Operations list is effectively identical to the input.

So enabling EnableLoopUnrolling today adds a traversal but does not actually unroll or transform the IR. That’s fine as a scaffold, but the XML docs and Name (“Loop Unrolling”) suggest a fully implemented optimization.

Two suggestions:

• Either implement the real unrolling / fusion behavior, or
• Explicitly mark this as a placeholder (e.g., doc comment and/or option name) so users don’t assume they’re getting loop-unrolling benefits yet.

Optionally, you might:

• Clone InputIds / OutputIds lists (new List<int>(graph.InputIds)) to avoid aliasing between original and optimized graphs.
• Remove or use MAX_OPS_TO_UNROLL to match the documented heuristics.

Also applies to: 146-246

src/JitCompiler/Optimizations/AutoTuningPass.cs (1)

121-127: Reuse shared shape utilities instead of manual Aggregate for tensor sizes.

You repeatedly compute tensor “sizes” via shape.Aggregate(1, (a, b) => a * b). Now that TensorShapeExtensions exists, this should likely use shape.GetElementCount() for consistency (and to encode the “dynamic dimension” semantics in one place). That will also avoid re‑implementing size logic if behavior for dynamic/empty shapes changes later.

Also applies to: 142-145, 190-200

src/JitCompiler/Optimizations/ConstantFoldingPass.cs (1)

79-142: Current implementation is safe but only performs foldability analysis, not actual folding.

Optimize conservatively keeps all operations and only annotates foldable ones via Metadata["FoldableOps"], with constant evaluation left as a future step (per comments). That’s a reasonable staged approach and shouldn’t change graph semantics today.

If you intend this pass to be “analysis‑only” for now, you might want to reflect that in the XML summary/remarks to avoid confusion about it already replacing ops with constants.

src/JitCompiler/IR/TensorShape.cs (1)

300-312: Shape validation rule matches comments; consider reusing it before broadcast operations.

IsValidShape correctly enforces “positive or -1, no zeros,” which matches the remarks. You may want to call this (or assert it) in places that rely on well‑formed shapes (e.g., before broadcasting) to fail fast on invalid shapes rather than propagating nonsense dimensions.

src/PredictionModelBuilder.cs (1)

67-67: JIT configuration and compilation flow look solid; consider null‑guarding CompilerOptions

The JIT wiring in ConfigureJitCompilation and BuildAsync is coherent: config defaults to Enabled = true when omitted, ThrowOnFailure semantics are respected, and unsupported/non‑JIT models fall back gracefully with clear console messages. Passing the compiled delegate into PredictionModelResult matches the usage pattern in PredictionModelResult.Predict.

The only robustness gap is assuming _jitCompilationConfig.CompilerOptions is always non‑null. Because CompilerOptions has a public setter, callers can explicitly set it to null and trigger a NullReferenceException in:

var jitCompiler = new AiDotNet.JitCompiler.JitCompiler(_jitCompilationConfig.CompilerOptions);

You can harden this by normalizing to a default instance:

- var jitCompiler = new AiDotNet.JitCompiler.JitCompiler(_jitCompilationConfig.CompilerOptions);
+ var options = _jitCompilationConfig.CompilerOptions ?? new JitCompilerOptions();
+ var jitCompiler = new AiDotNet.JitCompiler.JitCompiler(options);

This keeps the external surface flexible while avoiding a surprising crash on misconfigured inputs.

Also applies to: 269-338, 652-693, 709-710

docs/JIT-Compilation-Plan-Gap-Analysis.md (1)

636-799: Tighten wording around type‑safety bullets and align document version/status

The plan doc is very thorough; a couple of small edits would reduce confusion:

• In the “Challenge 2: Type Safety” bullet list, the line:

  - Runtime type checking where needed - Validated at compilation time

  reads as a fused sentence. Consider splitting into two bullets or rephrasing, e.g.:

  - Runtime type checking where needed  
  - Compile‑time validation of IR invariants

• The header declares Document Version: 3.0 – MAJOR UPDATE, while the history section later calls “Version 4.0 (Implementation Complete) ← CURRENT”. It would help future readers if the top‑level version/status reflects the current state (e.g., bump the header to 4.0 and adjust the “Status” line to match the “Implementation Complete / Ready for testing and integration” wording).

Both are editorial, but they make the long doc easier to trust at a glance.

src/JitCompiler/IR/Operations/FusedOps.cs (1)

27-230: Fused op definitions and validations align well with fusion patterns

The fused op classes look consistent with the fusion rules in OperationFusionPass and the tests:

• InputIds.Length checks correctly encode expected inputs (e.g., 3 for linear/ dense, 6 for Conv+BN, 2 for elementwise/residual).
• Guarding against empty ActivationName / ElementwiseOp in the activation‑fused ops is a good sanity check on IR integrity.

If you ever see invalid activation/elementwise names in practice, you might later tighten validation (e.g., compare against a known set or an enum), but what’s here is perfectly reasonable for v1.

src/JitCompiler/JitCompiler.cs (2)

50-54: Clarify or improve thread-safety of JitCompiler (shared builder/codegen).

JitCompiler holds a single _irBuilder and _codeGenerator instance that are reused across calls. Both IRBuilder and CodeGenerator maintain internal mutable state (e.g., _nextTensorId, _nodeToTensorId, _tensorVariables, _expressions), and the compile methods invoke them without synchronization.

If a single JitCompiler instance is used concurrently from multiple threads (e.g., compiling multiple graphs in parallel), this shared mutable state can lead to races, corrupted IR graphs, or incorrect compiled functions.

Options:

-    private readonly IRBuilder _irBuilder = new();
-    private readonly CodeGenerator _codeGenerator = new();
+    private readonly IRBuilder _irBuilder = new();
+    private readonly CodeGenerator _codeGenerator = new();
+
+    // Option A (simple): document JitCompiler as not thread-safe for Compile* methods.
+    // Option B (safer): guard compile paths with a lock, or
+    // Option C (more granular): use local IRBuilder/CodeGenerator instances per Compile* call.

At minimum, document the intended thread-safety guarantees; if you expect JitCompiler to be reused across threads, wrap calls to _irBuilder/_codeGenerator in a lock or refactor to per-call instances.

Also applies to: 185-215, 329-359, 384-422

---

424-455: Support classes (ApplyOptimizations, options, stats, cache stats) are well-structured.

• ApplyOptimizations cleanly composes passes in sequence.
• JitCompilerOptions is explicit about current vs planned features (loop unrolling, adaptive fusion, auto-tuning, SIMD hints).
• CompilationStats and CacheStats provide useful, human-readable diagnostics for profiling and debugging.

Only minor thought: OperationsEliminated/OptimizationPercentage can be negative if an optimization increases op count (e.g., some fusions), but that may be acceptable depending on how you interpret “optimization.”

Also applies to: 472-497, 514-587, 602-655, 668-688

src/JitCompiler/IR/Operations/BackwardOps.cs (1)

46-61: Backward gradient ops are consistent with IRBuilder and CodeGenerator usage.

The gradient IR ops are well-aligned with how IRBuilder.BuildBackward and CodeGenerator construct and consume them (input counts, ordering, and shapes). Validate() and ToString() implementations will be helpful for debugging backward graphs.

Two minor hardening ideas:

• For GradAddOp, GradSubtractOp, and GradElementwiseMultiplyOp, consider validating that InputIndex is in the expected range (0 or 1) to catch misconstructed IR earlier.
• For advanced grads (GradConv2DOp, GradMaxPool2DOp, GradBatchNormOp), once you wire them into IRBuilder.CreateBackwardOps, remember to add corresponding codegen paths; right now they’re safe placeholders.

Also applies to: 73-120, 131-149, 160-197, 208-245, 257-295, 305-345, 356-373, 385-401, 412-427

src/JitCompiler/IR/Operations/AllOtherOps.cs (1)

7-27: IR op definitions look consistent and validation-focused.

The forward IR ops here (reductions, shape ops, convolutions, pooling, normalization, and advanced ops) have coherent Validate() implementations and, where present, helpful ToString() overrides. The input-count checks align with how IRBuilder.ConvertNodeToOp constructs these operations, and basic parameter constraints (e.g., positive scales, bias handling, pooling inputs) are sensible.

If you want to tighten things further later, you could add shape/parameter sanity checks (e.g., matching lengths of Stride, Padding, OutputPadding, PoolSize) and ToString overrides for the remaining ops, but the current state is already quite usable.

Also applies to: 45-88, 94-134, 140-195, 202-225, 231-291, 297-329, 335-367, 373-431

📜 Review details

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 76b1c2d and 2371f17.

📒 Files selected for processing (41)

• docs/JIT-Compilation-Plan-Gap-Analysis.md (1 hunks)
• docs/JIT-Compiler-Implementation-Summary.md (1 hunks)
• docs/JIT-Compiler-Usage-Guide.md (1 hunks)
• docs/JIT-INTEGRATION-SUMMARY.md (1 hunks)
• examples/JitCompiler/BasicUsageExample.cs (1 hunks)
• examples/JitCompiler/README.md (1 hunks)
• src/Autodiff/ComputationNode.cs (1 hunks)
• src/Configuration/JitCompilationConfig.cs (1 hunks)
• src/Interfaces/IJitCompilable.cs (1 hunks)
• src/JitCompiler/CodeGen/CodeGenerator.cs (1 hunks)
• src/JitCompiler/CodeGen/GradientOps.cs (1 hunks)
• src/JitCompiler/CodeGen/SIMDOptimizer.cs (1 hunks)
• src/JitCompiler/IR/IRGraph.cs (1 hunks)
• src/JitCompiler/IR/IROp.cs (1 hunks)
• src/JitCompiler/IR/IRType.cs (1 hunks)
• src/JitCompiler/IR/Operations/ActivationOps.cs (1 hunks)
• src/JitCompiler/IR/Operations/AllOtherOps.cs (1 hunks)
• src/JitCompiler/IR/Operations/BackwardOps.cs (1 hunks)
• src/JitCompiler/IR/Operations/BasicArithmeticOps.cs (1 hunks)
• src/JitCompiler/IR/Operations/FusedOps.cs (1 hunks)
• src/JitCompiler/IR/Operations/MathOps.cs (1 hunks)
• src/JitCompiler/IR/Operations/MatrixOps.cs (1 hunks)
• src/JitCompiler/IR/TensorShape.cs (1 hunks)
• src/JitCompiler/IRBuilder.cs (1 hunks)
• src/JitCompiler/JitCompiler.cs (1 hunks)
• src/JitCompiler/Optimizations/AdaptiveFusionPass.cs (1 hunks)
• src/JitCompiler/Optimizations/AutoTuningPass.cs (1 hunks)
• src/JitCompiler/Optimizations/ConstantFoldingPass.cs (1 hunks)
• src/JitCompiler/Optimizations/DeadCodeEliminationPass.cs (1 hunks)
• src/JitCompiler/Optimizations/IOptimizationPass.cs (1 hunks)
• src/JitCompiler/Optimizations/LoopUnrollingPass.cs (1 hunks)
• src/JitCompiler/Optimizations/OperationFusionPass.cs (1 hunks)
• src/JitCompiler/README.md (1 hunks)
• src/Models/NeuralNetworkModel.cs (1 hunks)
• src/Models/Results/PredictionModelResult.cs (4 hunks)
• src/PredictionModelBuilder.cs (4 hunks)
• tests/AiDotNet.Tests/Benchmarks/JIT_BENCHMARKS_README.md (1 hunks)
• tests/AiDotNet.Tests/Benchmarks/JitCompilerBenchmarks.cs (1 hunks)
• tests/AiDotNet.Tests/UnitTests/JitCompiler/IRBuilderTests.cs (1 hunks)
• tests/AiDotNet.Tests/UnitTests/JitCompiler/JitCompilerTests.cs (1 hunks)
• tests/AiDotNet.Tests/UnitTests/JitCompiler/OptimizationPassTests.cs (1 hunks)

🧰 Additional context used

🧬 Code graph analysis (29)

src/JitCompiler/IR/Operations/MatrixOps.cs (2)

src/JitCompiler/IR/Operations/ActivationOps.cs (5)

• Validate (21-26)
• Validate (48-53)
• Validate (75-80)
• Validate (108-113)
• Validate (143-149)

src/JitCompiler/IR/Operations/AllOtherOps.cs (11)

• Validate (15-20)
• Validate (34-39)
• Validate (50-55)
• Validate (66-71)
• Validate (82-87)
• Validate (101-107)
• Validate (122-127)
• Validate (143-148)
• Validate (158-163)
• Validate (173-179)
• Validate (189-195)

src/Interfaces/IJitCompilable.cs (2)

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

• TOutput (626-663)

src/PredictionModelBuilder.cs (1)

• TOutput (729-732)

src/JitCompiler/IR/Operations/MathOps.cs (2)

src/JitCompiler/IR/Operations/ActivationOps.cs (5)

• Validate (21-26)
• Validate (48-53)
• Validate (75-80)
• Validate (108-113)
• Validate (143-149)

src/JitCompiler/IR/Operations/AllOtherOps.cs (11)

• Validate (15-20)
• Validate (34-39)
• Validate (50-55)
• Validate (66-71)
• Validate (82-87)
• Validate (101-107)
• Validate (122-127)
• Validate (143-148)
• Validate (158-163)
• Validate (173-179)
• Validate (189-195)

src/JitCompiler/CodeGen/CodeGenerator.cs (2)

src/JitCompiler/JitCompiler.cs (7)

• JitCompiler (48-498)
• JitCompiler (74-76)
• JitCompiler (99-137)
• Func (185-215)
• Func (239-277)
• Func (329-359)
• Func (384-422)

src/JitCompiler/CodeGen/GradientOps.cs (15)

• Tensor (33-45)
• Tensor (52-57)
• Tensor (64-76)
• Tensor (83-87)
• Tensor (94-99)
• Tensor (106-111)
• Tensor (118-124)
• Tensor (131-138)
• Tensor (145-152)
• Tensor (159-163)
• Tensor (170-174)
• Tensor (181-194)
• Tensor (199-213)
• Tensor (218-229)
• GradientOps (24-230)

src/Configuration/JitCompilationConfig.cs (1)

src/JitCompiler/JitCompiler.cs (1)

• JitCompilerOptions (514-587)

examples/JitCompiler/BasicUsageExample.cs (1)

src/JitCompiler/JitCompiler.cs (1)

• JitCompilerOptions (514-587)

src/JitCompiler/IR/IROp.cs (3)

src/JitCompiler/IR/IRType.cs (1)

• IRType (31-48)

src/JitCompiler/IR/IRGraph.cs (3)

• Validate (137-189)
• ToString (194-206)
• IRGraph (29-265)

src/JitCompiler/IR/TensorShape.cs (2)

• IsValidShape (300-312)
• ShapeToString (226-230)

src/JitCompiler/Optimizations/LoopUnrollingPass.cs (6)

src/JitCompiler/JitCompiler.cs (4)

• JitCompiler (48-498)
• JitCompiler (74-76)
• JitCompiler (99-137)
• IRGraph (445-455)

src/JitCompiler/IRBuilder.cs (5)

• IRGraph (69-112)
• IRGraph (509-621)
• List (398-423)
• List (630-794)
• IROp (145-315)

src/JitCompiler/CodeGen/SIMDOptimizer.cs (1)

• IsElementWiseOp (115-128)

src/JitCompiler/IR/Operations/BasicArithmeticOps.cs (5)

• AddOp (20-28)
• SubtractOp (44-52)
• ElementwiseMultiplyOp (71-79)
• DivideOp (95-103)
• NegateOp (153-161)

src/JitCompiler/IR/Operations/ActivationOps.cs (3)

• ReLUOp (19-27)
• SigmoidOp (46-54)
• TanhOp (73-81)

src/JitCompiler/IR/Operations/MathOps.cs (2)

• ExpOp (17-25)
• LogOp (41-49)

src/JitCompiler/Optimizations/AdaptiveFusionPass.cs (4)

src/JitCompiler/JitCompiler.cs (4)

• JitCompiler (48-498)
• JitCompiler (74-76)
• JitCompiler (99-137)
• IRGraph (445-455)

src/JitCompiler/IRBuilder.cs (5)

• IRGraph (69-112)
• IRGraph (509-621)
• List (398-423)
• List (630-794)
• IROp (145-315)

src/JitCompiler/Optimizations/AutoTuningPass.cs (2)

• IRGraph (84-103)
• IRGraph (208-218)

src/JitCompiler/Optimizations/OperationFusionPass.cs (1)

• OperationFusionPass (48-544)

src/JitCompiler/Optimizations/IOptimizationPass.cs (1)

src/JitCompiler/Optimizations/OperationFusionPass.cs (1)

• IRGraph (58-132)

src/JitCompiler/Optimizations/OperationFusionPass.cs (7)

src/JitCompiler/IRBuilder.cs (3)

• IRGraph (69-112)
• IRGraph (509-621)
• IROp (145-315)

src/JitCompiler/Optimizations/IOptimizationPass.cs (1)

• IRGraph (78-78)

src/JitCompiler/IR/Operations/MatrixOps.cs (1)

• MatMulOp (23-31)

src/JitCompiler/IR/Operations/BasicArithmeticOps.cs (4)

• AddOp (20-28)
• SubtractOp (44-52)
• ElementwiseMultiplyOp (71-79)
• DivideOp (95-103)

src/JitCompiler/IR/Operations/FusedOps.cs (6)

• FusedLinearOp (27-38)
• FusedLinearActivationOp (56-73)
• FusedDenseLayerOp (178-195)
• FusedElementwiseActivationOp (138-160)
• FusedConvBatchNormOp (90-121)
• FusedResidualBlockOp (213-230)

src/JitCompiler/IR/Operations/ActivationOps.cs (3)

• ReLUOp (19-27)
• SigmoidOp (46-54)
• TanhOp (73-81)

src/JitCompiler/IR/Operations/AllOtherOps.cs (2)

• Conv2DOp (205-225)
• BatchNormOp (355-367)

tests/AiDotNet.Tests/UnitTests/JitCompiler/JitCompilerTests.cs (1)

src/JitCompiler/JitCompiler.cs (2)

• JitCompilerOptions (514-587)
• CompilationStats (602-655)

tests/AiDotNet.Tests/UnitTests/JitCompiler/IRBuilderTests.cs (5)

src/JitCompiler/IRBuilder.cs (1)

• IRBuilder (32-795)

src/JitCompiler/IR/Operations/BasicArithmeticOps.cs (2)

• AddOp (20-28)
• PowerOp (119-137)

src/JitCompiler/IR/Operations/MatrixOps.cs (1)

• MatMulOp (23-31)

src/JitCompiler/IR/Operations/MathOps.cs (2)

• ExpOp (17-25)
• LogOp (41-49)

src/JitCompiler/IR/Operations/ActivationOps.cs (1)

• ReLUOp (19-27)

src/JitCompiler/Optimizations/AutoTuningPass.cs (3)

src/JitCompiler/JitCompiler.cs (4)

• JitCompiler (48-498)
• JitCompiler (74-76)
• JitCompiler (99-137)
• IRGraph (445-455)

src/JitCompiler/IRBuilder.cs (2)

• IRGraph (69-112)
• IRGraph (509-621)

src/JitCompiler/Optimizations/AdaptiveFusionPass.cs (3)

• IRGraph (86-109)
• IRGraph (153-195)
• IRGraph (200-205)

src/JitCompiler/IR/Operations/ActivationOps.cs (2)

src/JitCompiler/IR/Operations/AllOtherOps.cs (20)

• Validate (15-20)
• Validate (34-39)
• Validate (50-55)
• Validate (66-71)
• Validate (82-87)
• Validate (101-107)
• Validate (122-127)
• Validate (143-148)
• Validate (158-163)
• Validate (173-179)
• Validate (189-195)
• Validate (211-218)
• Validate (236-241)
• Validate (252-257)
• Validate (269-274)
• Validate (285-290)
• ToString (22-26)
• ToString (109-112)
• ToString (129-133)
• ToString (220-224)

src/JitCompiler/IR/TensorShape.cs (1)

• ShapeToString (226-230)

src/Models/Results/PredictionModelResult.cs (3)

src/Models/NeuralNetworkModel.cs (1)

• Tensor (547-554)

src/NeuralNetworks/NeuralNetworkBase.cs (9)

• Tensor (249-270)
• Tensor (286-289)
• Tensor (365-385)
• Tensor (393-400)
• Tensor (420-442)
• Tensor (479-524)
• Tensor (872-872)
• Tensor (1162-1184)
• T (1055-1064)

src/Models/VectorModel.cs (1)

• T (381-400)

src/JitCompiler/IRBuilder.cs (7)

src/JitCompiler/JitCompiler.cs (4)

• JitCompiler (48-498)
• JitCompiler (74-76)
• JitCompiler (99-137)
• IRGraph (445-455)

src/Autodiff/ComputationNode.cs (3)

• ComputationNode (28-414)
• ComputationNode (212-225)
• List (301-342)

src/JitCompiler/IR/Operations/BasicArithmeticOps.cs (6)

• AddOp (20-28)
• SubtractOp (44-52)
• ElementwiseMultiplyOp (71-79)
• DivideOp (95-103)
• PowerOp (119-137)
• NegateOp (153-161)

src/JitCompiler/IR/Operations/ActivationOps.cs (5)

• ReLUOp (19-27)
• SigmoidOp (46-54)
• TanhOp (73-81)
• SoftmaxOp (101-119)
• ApplyActivationOp (136-155)

src/JitCompiler/IR/Operations/MatrixOps.cs (2)

• MatMulOp (23-31)
• TransposeOp (53-61)

src/JitCompiler/IR/Operations/AllOtherOps.cs (18)

• ReshapeOp (97-113)
• ConcatOp (118-134)
• PadOp (139-149)
• CropOp (154-164)
• UpsampleOp (169-180)
• PixelShuffleOp (185-196)
• Conv2DOp (205-225)
• ConvTranspose2DOp (230-242)
• DepthwiseConv2DOp (247-258)
• DilatedConv2DOp (263-275)
• LocallyConnectedConv2DOp (280-291)
• MaxPool2DOp (300-312)
• AvgPool2DOp (317-329)
• LayerNormOp (338-350)
• BatchNormOp (355-367)
• AffineGridOp (390-400)
• GridSampleOp (405-416)
• RBFKernelOp (421-431)

src/JitCompiler/IR/Operations/BackwardOps.cs (12)

• GradAccumulateOp (46-61)
• GradAddOp (73-91)
• GradSubtractOp (102-120)
• GradElementwiseMultiplyOp (131-149)
• GradMatMulLeftOp (160-173)
• GradMatMulRightOp (184-197)
• GradReLUOp (208-221)
• GradSigmoidOp (232-245)
• GradTanhOp (256-269)
• GradExpOp (281-294)
• GradLogOp (305-318)
• GradSoftmaxOp (330-345)

src/JitCompiler/CodeGen/SIMDOptimizer.cs (3)

src/JitCompiler/JitCompiler.cs (3)

• JitCompiler (48-498)
• JitCompiler (74-76)
• JitCompiler (99-137)

src/JitCompiler/IRBuilder.cs (1)

• IROp (145-315)

src/JitCompiler/Optimizations/LoopUnrollingPass.cs (1)

• IsElementWiseOp (193-205)

src/JitCompiler/IR/TensorShape.cs (1)

src/JitCompiler/IR/IROp.cs (1)

• ToString (194-198)

src/JitCompiler/Optimizations/DeadCodeEliminationPass.cs (3)

src/JitCompiler/IRBuilder.cs (4)

• IRGraph (69-112)
• IRGraph (509-621)
• List (398-423)
• List (630-794)

src/JitCompiler/Optimizations/OperationFusionPass.cs (2)

• IRGraph (58-132)
• List (504-543)

src/JitCompiler/Optimizations/ConstantFoldingPass.cs (1)

• IRGraph (79-142)

src/JitCompiler/Optimizations/ConstantFoldingPass.cs (7)

src/JitCompiler/IRBuilder.cs (5)

• IRGraph (69-112)
• IRGraph (509-621)
• List (398-423)
• List (630-794)
• IROp (145-315)

src/JitCompiler/Optimizations/OperationFusionPass.cs (2)

• IRGraph (58-132)
• List (504-543)

src/JitCompiler/Optimizations/DeadCodeEliminationPass.cs (2)

• IRGraph (80-164)
• HashSet (187-231)

src/JitCompiler/IR/Operations/BasicArithmeticOps.cs (6)

• AddOp (20-28)
• SubtractOp (44-52)
• ElementwiseMultiplyOp (71-79)
• DivideOp (95-103)
• PowerOp (119-137)
• NegateOp (153-161)

src/JitCompiler/IR/Operations/ActivationOps.cs (4)

• ReLUOp (19-27)
• SigmoidOp (46-54)
• TanhOp (73-81)
• SoftmaxOp (101-119)

src/JitCompiler/IR/Operations/MatrixOps.cs (2)

• MatMulOp (23-31)
• TransposeOp (53-61)

src/JitCompiler/IR/Operations/AllOtherOps.cs (14)

• SumOp (10-27)
• MeanOp (32-40)
• ReduceMaxOp (45-56)
• ReduceMeanOp (61-72)
• ReduceLogVarianceOp (77-88)
• ReshapeOp (97-113)
• ConcatOp (118-134)
• PadOp (139-149)
• CropOp (154-164)
• Conv2DOp (205-225)
• MaxPool2DOp (300-312)
• AvgPool2DOp (317-329)
• LayerNormOp (338-350)
• BatchNormOp (355-367)

tests/AiDotNet.Tests/UnitTests/JitCompiler/OptimizationPassTests.cs (10)

tests/AiDotNet.Tests/UnitTests/JitCompiler/IRBuilderTests.cs (7)

• Fact (14-45)
• Fact (47-91)
• Fact (93-130)
• Fact (132-162)
• Fact (164-204)
• Fact (206-227)
• Fact (229-292)

src/JitCompiler/IR/Operations/BasicArithmeticOps.cs (2)

• AddOp (20-28)
• ElementwiseMultiplyOp (71-79)

src/JitCompiler/Optimizations/DeadCodeEliminationPass.cs (1)

• DeadCodeEliminationPass (38-258)

src/JitCompiler/IR/Operations/ActivationOps.cs (2)

• ReLUOp (19-27)
• SigmoidOp (46-54)

src/JitCompiler/IR/Operations/MathOps.cs (2)

• ExpOp (17-25)
• LogOp (41-49)

src/JitCompiler/IR/Operations/MatrixOps.cs (1)

• MatMulOp (23-31)

src/JitCompiler/Optimizations/OperationFusionPass.cs (1)

• OperationFusionPass (48-544)

src/JitCompiler/IR/Operations/FusedOps.cs (4)

• FusedLinearOp (27-38)
• FusedDenseLayerOp (178-195)
• FusedElementwiseActivationOp (138-160)
• FusedConvBatchNormOp (90-121)

src/JitCompiler/IR/Operations/AllOtherOps.cs (2)

• Conv2DOp (205-225)
• BatchNormOp (355-367)

src/JitCompiler/Optimizations/ConstantFoldingPass.cs (1)

• ConstantFoldingPass (38-269)

src/JitCompiler/JitCompiler.cs (4)

src/JitCompiler/IRBuilder.cs (5)

• IRBuilder (32-795)
• List (398-423)
• List (630-794)
• IRGraph (69-112)
• IRGraph (509-621)

src/JitCompiler/CodeGen/CodeGenerator.cs (3)

• CodeGenerator (50-565)
• CodeGenerator (73-79)
• Func (111-164)

src/JitCompiler/CodeGen/GradientOps.cs (14)

• Tensor (33-45)
• Tensor (52-57)
• Tensor (64-76)
• Tensor (83-87)
• Tensor (94-99)
• Tensor (106-111)
• Tensor (118-124)
• Tensor (131-138)
• Tensor (145-152)
• Tensor (159-163)
• Tensor (170-174)
• Tensor (181-194)
• Tensor (199-213)
• Tensor (218-229)

src/JitCompiler/IR/IRGraph.cs (1)

• ComputeStructureHash (229-264)

src/JitCompiler/CodeGen/GradientOps.cs (1)

src/Autodiff/TensorOperations.cs (1)

• TensorOperations (45-5389)

src/JitCompiler/IR/IRGraph.cs (2)

src/JitCompiler/IR/IROp.cs (3)

• IROp (29-199)
• Validate (162-173)
• ToString (194-198)

src/JitCompiler/IR/TensorShape.cs (1)

• GetShapeHashCode (249-257)

src/JitCompiler/IR/Operations/BackwardOps.cs (4)

src/JitCompiler/JitCompiler.cs (5)

• JitCompiler (48-498)
• JitCompiler (74-76)
• JitCompiler (99-137)
• ToString (645-654)
• ToString (683-688)

src/JitCompiler/IRBuilder.cs (1)

• IROp (145-315)

src/JitCompiler/Optimizations/AdaptiveFusionPass.cs (2)

• IROp (253-257)
• IROp (272-277)

src/JitCompiler/IR/TensorShape.cs (1)

• ShapeToString (226-230)

src/JitCompiler/IR/Operations/AllOtherOps.cs (1)

src/JitCompiler/IR/TensorShape.cs (1)

• ShapeToString (226-230)

src/PredictionModelBuilder.cs (3)

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

• TOutput (626-663)

src/Configuration/JitCompilationConfig.cs (1)

• JitCompilationConfig (37-141)

src/Interfaces/IJitCompilable.cs (1)

• ComputationNode (84-84)

src/JitCompiler/IR/Operations/BasicArithmeticOps.cs (3)

src/JitCompiler/IR/Operations/ActivationOps.cs (7)

• Validate (21-26)
• Validate (48-53)
• Validate (75-80)
• Validate (108-113)
• Validate (143-149)
• ToString (115-118)
• ToString (151-154)

src/JitCompiler/IR/Operations/AllOtherOps.cs (15)

• Validate (15-20)
• Validate (34-39)
• Validate (50-55)
• Validate (66-71)
• Validate (82-87)
• Validate (101-107)
• Validate (122-127)
• Validate (143-148)
• Validate (158-163)
• Validate (173-179)
• Validate (189-195)
• ToString (22-26)
• ToString (109-112)
• ToString (129-133)
• ToString (220-224)

src/JitCompiler/IR/TensorShape.cs (1)

• ShapeToString (226-230)

🪛 GitHub Actions: Build

src/JitCompiler/JitCompiler.cs

[error] 340-340: CS1002: ; expected

🪛 GitHub Actions: Quality Gates (.NET)

src/JitCompiler/JitCompiler.cs

[error] 340-340: CS1002: ; expected

🪛 GitHub Check: Build All Frameworks

src/JitCompiler/JitCompiler.cs

[failure] 395-395:
; expected

---

[failure] 340-340:
; expected

---

[failure] 395-395:
; expected

---

[failure] 340-340:
; expected

---

[failure] 395-395:
; expected

---

[failure] 340-340:
; expected

---

[failure] 395-395:
; expected

---

[failure] 340-340:
; expected

🪛 GitHub Check: Publish Size Analysis

src/JitCompiler/JitCompiler.cs

[failure] 395-395:
; expected

---

[failure] 340-340:
; expected

🪛 LanguageTool

examples/JitCompiler/README.md

[style] ~188-~188: As an alternative to the over-used intensifier ‘very’, consider replacing this phrase.
Context: ...aphs that change structure frequently - Very small operations (compilation overhead) ### ...

(EN_WEAK_ADJECTIVE)

docs/JIT-Compiler-Usage-Guide.md

[style] ~211-~211: As an alternative to the over-used intensifier ‘very’, consider replacing this phrase.
Context: ...aphs that change structure frequently - Very small operations (compilation overhead) ## C...

(EN_WEAK_ADJECTIVE)

docs/JIT-Compiler-Implementation-Summary.md

[uncategorized] ~479-~479: If this is a compound adjective that modifies the following noun, use a hyphen.
Context: ...ence - Zero Breaking Changes: Fully backward compatible - Comprehensive Testing: 20+ unit t...

(EN_COMPOUND_ADJECTIVE_INTERNAL)

docs/JIT-Compilation-Plan-Gap-Analysis.md

[style] ~755-~755: The double modal “needed Validated” is nonstandard (only accepted in certain dialects). Consider “to be Validated”.
Context: ... - Runtime type checking where needed - Validated at compilation time ### Challenge 3: D...

(NEEDS_FIXED)

---

[grammar] ~777-~777: Ensure spelling is correct
Context: ...round) - Compilation budget (abort if > 100ms for simple graphs) --- ## Success Met...

(QB_NEW_EN_ORTHOGRAPHY_ERROR_IDS_1)

🪛 markdownlint-cli2 (0.18.1)

examples/JitCompiler/README.md

83-83: Fenced code blocks should have a language specified

(MD040, fenced-code-language)

---

96-96: Fenced code blocks should have a language specified

(MD040, fenced-code-language)

---

111-111: Fenced code blocks should have a language specified

(MD040, fenced-code-language)

---

124-124: Fenced code blocks should have a language specified

(MD040, fenced-code-language)

---

145-145: Fenced code blocks should have a language specified

(MD040, fenced-code-language)

tests/AiDotNet.Tests/Benchmarks/JIT_BENCHMARKS_README.md

111-111: Fenced code blocks should have a language specified

(MD040, fenced-code-language)

---

287-287: Fenced code blocks should have a language specified

(MD040, fenced-code-language)

src/JitCompiler/README.md

29-29: Fenced code blocks should have a language specified

(MD040, fenced-code-language)

---

50-50: Fenced code blocks should have a language specified

(MD040, fenced-code-language)

---

105-105: Fenced code blocks should have a language specified

(MD040, fenced-code-language)

---

111-111: Fenced code blocks should have a language specified

(MD040, fenced-code-language)

---

117-117: Fenced code blocks should have a language specified

(MD040, fenced-code-language)

docs/JIT-Compiler-Implementation-Summary.md

258-258: Emphasis used instead of a heading

(MD036, no-emphasis-as-heading)

---

273-273: Emphasis used instead of a heading

(MD036, no-emphasis-as-heading)

---

286-286: Emphasis used instead of a heading

(MD036, no-emphasis-as-heading)

---

296-296: Emphasis used instead of a heading

(MD036, no-emphasis-as-heading)

---

459-459: Emphasis used instead of a heading

(MD036, no-emphasis-as-heading)

---

474-474: Emphasis used instead of a heading

(MD036, no-emphasis-as-heading)

---

499-499: Emphasis used instead of a heading

(MD036, no-emphasis-as-heading)

docs/JIT-INTEGRATION-SUMMARY.md

83-83: Fenced code blocks should have a language specified

(MD040, fenced-code-language)

---

230-230: Multiple headings with the same content

(MD024, no-duplicate-heading)

---

387-387: Fenced code blocks should have a language specified

(MD040, fenced-code-language)

⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)

• GitHub Check: CodeQL analysis (csharp)

🔇 Additional comments (22)

src/Autodiff/ComputationNode.cs (1)

136-161: LGTM - Well-documented property addition.

The OperationType property is well-designed for JIT compilation metadata. The nullable type ensures backward compatibility, and the XML documentation is comprehensive with beginner-friendly explanations.

src/JitCompiler/README.md (1)

29-46: Static analysis hints can be safely ignored.

The markdown linter flagged code blocks without language specifiers (lines 29, 50, 105, 111, 117). These blocks contain ASCII diagrams and transformation examples, not executable code, so language specifiers would be inappropriate here. The current formatting is correct and clear.

Also applies to: 50-75, 105-119

docs/JIT-Compiler-Implementation-Summary.md (1)

1-515: Comprehensive implementation documentation.

This is an excellent summary document that clearly captures the JIT compiler implementation scope, architecture, testing strategy, and future work. The level of detail is appropriate for both developers and stakeholders.

tests/AiDotNet.Tests/Benchmarks/JitCompilerBenchmarks.cs (1)

231-242: Good use of fixed random seed for reproducibility.

Using a fixed seed (42) in CreateRandomTensor ensures benchmark results are consistent across runs, which is essential for reliable performance comparisons.

tests/AiDotNet.Tests/Benchmarks/JIT_BENCHMARKS_README.md (1)

1-310: Comprehensive benchmark documentation.

This README provides excellent guidance for running and interpreting JIT compiler benchmarks. The troubleshooting section and best practices are particularly valuable.

The static analysis hints about missing language specifiers (lines 111, 287) can be ignored—these blocks contain expected output/pseudo-code rather than executable code.

src/Models/NeuralNetworkModel.cs (1)

23-49: Well-documented future enhancement plan.

The TODO comment provides clear, actionable guidance for implementing JIT compilation support for neural networks. The implementation approach and example usage are helpful for future development. This is an appropriate way to document planned enhancements without introducing incomplete code.

examples/JitCompiler/README.md (1)

1-262: Excellent examples documentation with clear learning path.

This README provides comprehensive guidance for JIT compiler examples, including:

• Clear overview of 5 example scenarios
• Multiple options for running examples
• Expected output with concrete examples
• Best practices and common issues

The static analysis hints about missing language specifiers (lines 83-145) can be ignored—these blocks show expected console output rather than code to be executed.

src/JitCompiler/IR/Operations/MathOps.cs (1)

3-73: Unary math IR ops and validation look consistent with existing patterns.

Docs are clear, and Validate() matches the existing unary-ops pattern (base.Validate + single input). No issues from a correctness or API‑shape perspective.

src/JitCompiler/IR/Operations/MatrixOps.cs (1)

3-61: Matrix op IR definitions and validation look correct.

MatMulOp and TransposeOp follow the same validation style as other ops, with appropriate input-arity checks and clear documentation about semantics and corresponding tensor operations. No issues from this file.

src/JitCompiler/CodeGen/CodeGenerator.cs (1)

278-349: Incorrect parameter type hints in FindMethod calls confirm real type mismatch

The review comment is correct: the Generate*Op methods pass typeof(ComputationNode<T>) to FindMethod, but _tensorVariables stores ParameterExpression objects created as Expression.Variable(typeof(Tensor<T>), ...) (line 122). The actual inputs are Tensor<T> variables, not ComputationNode<T>.

Additionally, the FindMethod implementation (lines 433–437) validates only the method name and parameter count; it ignores the parameter type hints entirely:

var method = _tensorOperationsMethods.FirstOrDefault(m =>
    m.Name == methodName &&
    m.GetParameters().Length == parameterTypes.Length);  // ← Types not checked

Meanwhile, all TensorOperations methods expect ComputationNode<T> parameters (e.g., public static ComputationNode<T> Add(ComputationNode<T> a, ComputationNode<T> b)), not Tensor<T>.

This creates a type mismatch that passes parameter count validation but could silently bind to the wrong overload if multiple signatures exist with the same parameter count. The suggestion to pass actual types (typeof(Tensor<T>)) and validate them in FindMethod is essential for robustness.

tests/AiDotNet.Tests/UnitTests/JitCompiler/IRBuilderTests.cs (1)

14-45: IRBuilder test coverage looks strong and well‑targeted

The tests collectively exercise core behaviors: basic op mapping, linear chains, DAGs with shared nodes, parameterized ops, and error handling for missing OperationType, plus a minimal but correct topological‑order assertion. They read cleanly and align with the described IRBuilder semantics.

Also applies to: 48-91, 93-203, 229-292

src/JitCompiler/IR/IROp.cs (2)

29-199: IROp base class and validation look coherent with IRGraph usage

The base IROp contract (output id, inputs, type, shape, OpType, basic Validate, ToString) is clean and aligns with how IRGraph.Validate and hashing consume it. Input-specific arity checks are appropriately delegated to derived ops.

No changes requested here.

---

201-280: IOptimizationPass surface is minimal and fits the optimization pipeline

IOptimizationPass with IRGraph Optimize(IRGraph) and string Name { get; } matches the usage pattern in JitCompiler.ApplyOptimizations and is easy to extend with new passes.

Interface definition looks good as-is.

src/JitCompiler/IR/IRGraph.cs (2)

29-189: IRGraph structure and validation are consistent with intended IR semantics

The IRGraph layout and Validate() implementation:

• Enforce that all inputs have shapes, operations are individually valid, inputs are produced before being consumed, and outputs correspond to produced tensors.
• Opportunistically fill in missing TensorShapes entries from each op’s OutputShape.

This matches the expectations of the rest of the JIT stack and should catch most structural issues early.

Looks solid.

---

191-264: Structure hash design looks appropriate for graph caching

ComputeStructureHash() hashing inputs (including shapes), operations (OpType, ids, types, shapes, inputs), and outputs provides a reasonably strong fingerprint for graph structure while ignoring actual tensor values.

This should work well for cache keys for compiled graphs given the current IR design.

src/JitCompiler/IR/Operations/BasicArithmeticOps.cs (1)

20-161: Basic arithmetic ops align with IRBuilder expectations

The arithmetic ops:

• Correctly extend IROp and override Validate() to enforce input arity (2 for binary ops, 1 for unary ops).
• Provide clear XML docs, and PowerOp exposes Exponent plus a helpful ToString() override.

These definitions match how IRBuilder.ConvertNodeToOp constructs them and should integrate cleanly with validation and codegen.

No changes requested.

src/JitCompiler/Optimizations/OperationFusionPass.cs (1)

501-543: Fusion opportunity introspection API looks good.

IdentifyFusionOpportunities is a nice, cheap analysis helper that mirrors the implemented fusion patterns and doesn’t mutate the graph. The pattern strings are clear and should be useful for debugging and tooling.

src/JitCompiler/Optimizations/ConstantFoldingPass.cs (1)

172-228: Foldability classification looks consistent with IR semantics.

The CanFold switch correctly limits folding to pure/deterministic operations (arithmetic, activations, shape ops, reductions, conv/pooling, normalization) and defaults to “not foldable” for unknown ops. This is a good conservative stance and aligns with the remarks about only folding operations known to be side‑effect‑free.

src/JitCompiler/Optimizations/AdaptiveFusionPass.cs (1)

86-109: Adaptive fusion control flow is coherent; conservative path is intentionally a stub.

The strategy selection and delegation to OperationFusionPass for Standard/Aggressive modes are straightforward and keep logic centralized. Given CreateFusedOp currently returns null, the Conservative path effectively leaves the graph unchanged, which matches the TODOs in the remarks about future fused op creation.

No functional changes here beyond metadata loss noted above.

Also applies to: 150-205, 270-277

src/JitCompiler/IR/TensorShape.cs (1)

103-125: Broadcasting and shape hashing utilities look solid and consistent.

The broadcasting compatibility check, result‑shape computation (including -1 as a dynamic dimension), and GetShapeHashCode implementation are all clear and should work well for JIT caching and shape reasoning.

Nice use of right‑to‑left alignment and dynamic propagation for broadcast semantics.

Also applies to: 150-179, 232-257

tests/AiDotNet.Tests/UnitTests/JitCompiler/OptimizationPassTests.cs (1)

15-392: Good, focused coverage of DCE, fusion patterns, and constant folding

These tests nicely pin down the intended behavior of the optimization passes:

• DCE tests exercise straight‑line graphs, a diamond pattern, and the statistics helper.
• Fusion tests cover the main fusion patterns (MatMul+Add, MatMul+Add+Activation, elementwise+activation, Conv+BN) plus the “multiple consumers” guard and fusion‑opportunity reporting.
• Constant folding tests are pragmatic smoke tests that ensure the pass runs and is wired correctly.

This suite should make future changes to the optimization passes much safer.

src/JitCompiler/JitCompiler.cs (1)

185-215: Overall forward compilation and stats logic looks solid.

Forward compilation flow is coherent: build IR → optional cache lookup → apply configured passes → codegen → cache store. CompileWithStats correctly tracks original vs optimized operation counts, records optimization pass names, and uses CacheHit to short-circuit compilation when appropriate.

Also applies to: 239-277

[coderabbitai]

⚠️ Potential issue | 🟠 Major

Fix JitCompiledModel example to match the public Compile signature

In the JitCompiledModel example, Compile is called with new[] { inputNode }, but the API reference below documents Compile<T>(ComputationNode<T> outputNode, List<ComputationNode<T>> inputs). Unless there is an overload taking an array, this snippet will not compile as-written.

Consider changing the example to pass a List<ComputationNode<T>> (or updating the API reference if an overload exists) so the docs and public surface stay in sync.

🤖 Prompt for AI Agents

In docs/JIT-Compiler-Usage-Guide.md around lines 218 to 233, the
JitCompiledModel example calls _jit.Compile(output, new[] { inputNode }) but the
public API expects a List<ComputationNode<T>>; update the example to pass a List
containing inputNode (e.g., new List<ComputationNode<float>> { inputNode }) so
it matches the documented Compile<T>(ComputationNode<T>,
List<ComputationNode<T>>) signature and compiles cleanly, or alternatively
adjust the API docs if an array overload actually exists.

[coderabbitai]

⚠️ Potential issue | 🟠 Major

CodeGenerator is not thread-safe due to shared mutable state

Generate<T> mutates the instance fields _tensorVariables and _expressions:

_tensorVariables.Clear();
_expressions.Clear();
// ...
_tensorVariables[inputId] = inputVar;
_expressions.Add(assignment);
// ...
_expressions.Add(outputArray);

JitCompiler holds a single CodeGenerator instance and may be used from multiple threads; concurrent calls to Compile / Generate will race on these collections, corrupting the generated expression tree and potentially throwing at runtime.

Consider one of:

• Making CodeGenerator.Generate<T> allocate local dictionaries/lists instead of using instance fields; or
• Making CodeGenerator internal to a per‑compilation scope instance; or
• Synchronizing access (e.g., a lock) around Generate<T> in JitCompiler if you want to keep a shared instance.

Given JIT compilation is often used in multi‑threaded inference scenarios, keeping this path thread‑safe is important.

[coderabbitai]

⚠️ Potential issue | 🔴 Critical

Bug: FindMethod generic handling will throw for Concat (array parameter)

GenerateConcatOp calls:

var method = FindMethod("Concat", typeof(ComputationNode<T>[]), typeof(int));

Inside FindMethod you do:

if (method.IsGenericMethodDefinition)
{
    var genericArg = parameterTypes[0].GetGenericArguments()[0];
    method = method.MakeGenericMethod(genericArg);
}

When parameterTypes[0] is an array type (ComputationNode<T>[]), GetGenericArguments() returns an empty array, so indexing [0] will throw IndexOutOfRangeException at code‑generation time for any IR graph that uses Concat.

A more robust approach is to use typeof(T) directly (or the method’s generic arguments) instead of inferring from parameterTypes[0], and to avoid assuming parameterTypes[0] is a generic type.

One possible fix:

-    private MethodInfo FindMethod(string methodName, params Type[] parameterTypes)
+    private MethodInfo FindMethod<T>(string methodName, params Type[] parameterTypes)
     {
         var method = _tensorOperationsMethods.FirstOrDefault(m =>
             m.Name == methodName &&
             m.GetParameters().Length == parameterTypes.Length);

         if (method == null)
         {
             throw new InvalidOperationException(
                 $"Could not find TensorOperations method '{methodName}' with {parameterTypes.Length} parameters");
         }

-        // If method is generic, make it concrete with T
-        if (method.IsGenericMethodDefinition)
-        {
-            var genericArg = parameterTypes[0].GetGenericArguments()[0];
-            method = method.MakeGenericMethod(genericArg);
-        }
+        // If method is generic, specialize it with the element type T
+        if (method.IsGenericMethodDefinition)
+        {
+            method = method.MakeGenericMethod(typeof(T));
+        }

         return method;
     }

And update all call sites to FindMethod<T>(...) accordingly (e.g., FindMethod<T>("Concat", ...), FindMethod<T>("Add", ...), etc.). This removes the array/generic-arg coupling and prevents the runtime exception.

🤖 Prompt for AI Agents

In src/JitCompiler/CodeGen/CodeGenerator.cs around lines 323 to 353,
GenerateConcatOp currently calls FindMethod("Concat",
typeof(ComputationNode<T>[]), typeof(int)) which leads FindMethod to assume
parameterTypes[0] is a generic type and call GetGenericArguments()[0], causing
an IndexOutOfRangeException for array types; fix by changing FindMethod to a
generic method FindMethod<TMethod>() (or equivalent) that uses the provided type
parameter (typeof(T) or the method’s generic parameter) to construct the generic
method instead of inferring from parameterTypes, update code paths to call
FindMethod<T>("Concat", typeof(ComputationNode<T>[]), typeof(int)) (and other
call sites like Add/Reduce) and also harden FindMethod so it checks for array
types and falls back safely (e.g., detect if parameterTypes[0].IsArray and use
parameterTypes[0].GetElementType() or the explicit generic T) before calling
MakeGenericMethod.

[coderabbitai]

⚠️ Potential issue | 🟠 Major

Fix generic‑type safety in helpers and clarify broadcasting assumptions for gradients

The per‑op gradient formulas look correct and line up with the existing TensorOperations semantics, but there are a couple of important implementation issues:

1. Generic type safety in CreateMask / CreateOnes

   Both helpers currently assume T can be treated as a double:

   dynamic val = inputData[i];
   resultData[i] = val > 0 ? (T)(object)1.0 : (T)(object)0.0;
   ...
   data[i] = (T)(object)1.0;

   This will throw at runtime for valid numeric types that aren’t directly castable from a boxed double (e.g., decimal, custom numeric types), and it bypasses the library’s MathHelper.GetNumericOperations<T>() abstraction.

   Since AiDotNet is generic over T, these helper methods should use the existing numeric operations instead of dynamic/boxing. For example (pseudocode):

   var numOps = MathHelper.GetNumericOperations<T>();
   
   // In CreateMask
   resultData[i] = numOps.GreaterThan(inputData[i], numOps.Zero)
       ? numOps.One
       : numOps.Zero;
   
   // In CreateOnes
   for (int i = 0; i < totalSize; i++)
       data[i] = numOps.One;

   This keeps gradients usable for all supported numeric types and matches how other parts of the autodiff stack handle T.

2. Broadcasting behavior for Add/Subtract gradients

   GradAdd and GradSubtract currently just return gradOutput (or its negation) and include a comment about broadcasting:

   // May need to handle broadcasting by summing over broadcasted dimensions
   return gradOutput;

   That’s only correct when the forward operation did not broadcast either input. If IR graphs can include broadcasted Add/Subtract, the backward pass should reduce over broadcasted dimensions (as TensorOperations.Sum does in the non‑JIT autodiff path), or the IR builder should explicitly avoid using these gradient helpers for broadcasted cases. It would be good either to implement the reduction logic here or to document/enforce the “no broadcasting” assumption at IR construction time.

3. Minor cleanup / perf

   • CreateMask constructs an unused Tensor<T> result then ignores it and returns a new Tensor<T> from resultData. You can drop the first allocation and just build from the array:

• var result = new Tensor(input.Shape);
• var inputData = input.ToArray();
• var resultData = result.ToArray();

• var inputData = input.ToArray();
• var resultData = new T[inputData.Length];

  
  - `shape.Aggregate(1, (a, b) => a * b)` in `CreateOnes` depends on LINQ; just make sure `System.Linq` is either explicitly imported or covered by your global usings.
  
  

Addressing (1) is the key correctness fix; (2) is important if you intend to JIT‑compile graphs that rely on broadcasting semantics.

🤖 Prompt for AI Agents

In src/JitCompiler/CodeGen/GradientOps.cs around lines 26 to 229, the helper
methods CreateMask and CreateOnes use dynamic/boxed double casts which are
unsafe for generic numeric types and an extra unused Tensor is allocated, and
GradAdd/GradSubtract ignore broadcasting; fix by using the library numeric
abstraction MathHelper.GetNumericOperations<T>() to obtain Zero/One and a
GreaterThan operation and replace dynamic/boxing with numOps.GreaterThan checks
and numOps.One/Zero assignments, fill the ones array using numOps.One rather
than casting from double, remove the unused Tensor allocation in CreateMask, and
for GradAdd/GradSubtract either implement reduction over broadcasted dimensions
(use TensorOperations.Sum with keepDims where inputs were broadcast) or
assert/document that inputs are same-shape and throw/handle accordingly so
gradients are correct when broadcasting occurs.

[coderabbitai]

⚠️ Potential issue | 🟠 Major

Align IRType enum with IRTypeExtensions mappings.

IRType exposes Half and Complex, but:

• FromSystemType never returns these values.
• ToSystemType throws NotSupportedException for them via the default branch.

That means any code that sets IRType.Half or IRType.Complex (e.g., when constructing IR nodes) will later fail if ToSystemType is used.

Consider one of these approaches:

• Add full support (if TFMs/refs allow it):

+using System.Numerics;

public static IRType FromSystemType(Type type)
{
    return type switch
    {
        ...
+       Type t when t == typeof(Half) => IRType.Half,
+       Type t when t == typeof(Complex) => IRType.Complex,
        _ => throw new NotSupportedException($"Type {type} not supported in IR")
    };
}

public static Type ToSystemType(this IRType irType)
{
    return irType switch
    {
        ...
+       IRType.Half => typeof(Half),
+       IRType.Complex => typeof(Complex),
        _ => throw new NotSupportedException($"IRType {irType} conversion not supported")
    };
}

• Or drop Half/Complex from the enum (or clearly mark them as unsupported), if you don’t plan to support these types yet.

Right now the API surface suggests support that the helpers don’t actually provide.

📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change

	public enum IRType
	{
	Float32,
	Float64,
	Int32,
	Int64,
	Byte,
	SByte,
	Int16,
	UInt16,
	UInt32,
	UInt64,
	Decimal,
	Half,
	Complex
	}
	/// <summary>
	/// Helper methods for IRType.
	/// </summary>
	public static class IRTypeExtensions
	{
	/// <summary>
	/// Gets the IRType for a given System.Type.
	/// </summary>
	public static IRType FromSystemType(Type type)
	{
	return type switch
	{
	Type t when t == typeof(float) => IRType.Float32,
	Type t when t == typeof(double) => IRType.Float64,
	Type t when t == typeof(int) => IRType.Int32,
	Type t when t == typeof(long) => IRType.Int64,
	Type t when t == typeof(byte) => IRType.Byte,
	Type t when t == typeof(sbyte) => IRType.SByte,
	Type t when t == typeof(short) => IRType.Int16,
	Type t when t == typeof(ushort) => IRType.UInt16,
	Type t when t == typeof(uint) => IRType.UInt32,
	Type t when t == typeof(ulong) => IRType.UInt64,
	Type t when t == typeof(decimal) => IRType.Decimal,
	_ => throw new NotSupportedException($"Type {type} not supported in IR")
	};
	}
	/// <summary>
	/// Gets the System.Type for a given IRType.
	/// </summary>
	public static Type ToSystemType(this IRType irType)
	{
	return irType switch
	{
	IRType.Float32 => typeof(float),
	IRType.Float64 => typeof(double),
	IRType.Int32 => typeof(int),
	IRType.Int64 => typeof(long),
	IRType.Byte => typeof(byte),
	IRType.SByte => typeof(sbyte),
	IRType.Int16 => typeof(short),
	IRType.UInt16 => typeof(ushort),
	IRType.UInt32 => typeof(uint),
	IRType.UInt64 => typeof(ulong),
	IRType.Decimal => typeof(decimal),
	_ => throw new NotSupportedException($"IRType {irType} conversion not supported")
	};
	}
	}
	using System.Numerics;
	public enum IRType
	{
	Float32,
	Float64,
	Int32,
	Int64,
	Byte,
	SByte,
	Int16,
	UInt16,
	UInt32,
	UInt64,
	Decimal,
	Half,
	Complex
	}
	/// <summary>
	/// Helper methods for IRType.
	/// </summary>
	public static class IRTypeExtensions
	{
	/// <summary>
	/// Gets the IRType for a given System.Type.
	/// </summary>
	public static IRType FromSystemType(Type type)
	{
	return type switch
	{
	Type t when t == typeof(float) => IRType.Float32,
	Type t when t == typeof(double) => IRType.Float64,
	Type t when t == typeof(int) => IRType.Int32,
	Type t when t == typeof(long) => IRType.Int64,
	Type t when t == typeof(byte) => IRType.Byte,
	Type t when t == typeof(sbyte) => IRType.SByte,
	Type t when t == typeof(short) => IRType.Int16,
	Type t when t == typeof(ushort) => IRType.UInt16,
	Type t when t == typeof(uint) => IRType.UInt32,
	Type t when t == typeof(ulong) => IRType.UInt64,
	Type t when t == typeof(decimal) => IRType.Decimal,
	Type t when t == typeof(Half) => IRType.Half,
	Type t when t == typeof(Complex) => IRType.Complex,
	_ => throw new NotSupportedException($"Type {type} not supported in IR")
	};
	}
	/// <summary>
	/// Gets the System.Type for a given IRType.
	/// </summary>
	public static Type ToSystemType(this IRType irType)
	{
	return irType switch
	{
	IRType.Float32 => typeof(float),
	IRType.Float64 => typeof(double),
	IRType.Int32 => typeof(int),
	IRType.Int64 => typeof(long),
	IRType.Byte => typeof(byte),
	IRType.SByte => typeof(sbyte),
	IRType.Int16 => typeof(short),
	IRType.UInt16 => typeof(ushort),
	IRType.UInt32 => typeof(uint),
	IRType.UInt64 => typeof(ulong),
	IRType.Decimal => typeof(decimal),
	IRType.Half => typeof(Half),
	IRType.Complex => typeof(Complex),
	_ => throw new NotSupportedException($"IRType {irType} conversion not supported")
	};
	}
	}

🤖 Prompt for AI Agents

In src/JitCompiler/IR/IRType.cs around lines 6 to 71, the IRType enum includes
Half and Complex but the extension methods neither return nor map those values;
fix by either (A) adding full mappings: in FromSystemType add cases for
typeof(System.Half) => IRType.Half and typeof(System.Numerics.Complex) =>
IRType.Complex (add using System.Numerics if needed and ensure target TFM
supports System.Half), and in ToSystemType return typeof(System.Half) and
typeof(System.Numerics.Complex) for IRType.Half/Complex respectively; or (B)
remove IRType.Half and IRType.Complex from the enum (and update any callers) or
mark them [Obsolete] with clear comments to indicate they are unsupported—choose
one approach and make corresponding unit-test/usage updates.

[coderabbitai]

⚠️ Potential issue | 🟠 Major

🧩 Analysis chain

Make tuning cache thread‑safe or clearly scope it per compilation.

_tuningCache is a mutable Dictionary<int, TuningConfig> shared by the AutoTuningPass instance. If a single JitCompiler (and its pass instances) is used concurrently across requests, parallel calls to Optimize can race on this dictionary (reads + writes, TryGetValue, indexer assignment), which is undefined and can throw or corrupt state.

Consider either:

• switching to ConcurrentDictionary<int, TuningConfig>, or
• scoping the cache per compilation (e.g., passing it in via IRGraph.Metadata or options) so each Optimize call works on isolated state.

If the design guarantees single‑threaded use of each JitCompiler, documenting that explicitly would also help.

---

I'll verify the thread-safety concern by investigating the codebase structure and usage patterns.

---

🏁 Script executed:

```shell
#!/bin/bash
# Search for JitCompiler class definition and instantiation patterns
rg -nP --type=cs 'class\s+JitCompiler\b' -A 10

Length of output: 852

---

🏁 Script executed:

#!/bin/bash
# Find where AutoTuningPass is instantiated
rg -nP --type=cs 'new\s+AutoTuningPass|AutoTuningPass\s*\(' -B 3 -A 3

Length of output: 370

---

🏁 Script executed:

#!/bin/bash
# Search for concurrent/threading patterns in the codebase
rg -nP --type=cs '(ConcurrentDictionary|lock\s*\(|Interlocked|Monitor\.|Thread|Task\.Run)' -C 2

Length of output: 88385

---

🏁 Script executed:

#!/bin/bash
# Check if JitCompiler is marked as thread-safe or has threading documentation
rg -nP --type=cs '(thread.safe|concurrent|parallel|async)' -i src/JitCompiler/ -C 2

Length of output: 1051

---

🏁 Script executed:

#!/bin/bash
# Search for Dictionary<int usage patterns in the codebase to see threading precedent
rg -nP --type=cs 'Dictionary<int' -B 2 -A 2

Length of output: 54833

---

---

🏁 Script executed:

```shell
#!/bin/bash
# Search for where Optimize is called in JitCompiler
rg -nP --type=cs 'Optimize\(' src/JitCompiler/ -B 2 -A 2

Length of output: 4237

---

🏁 Script executed:

#!/bin/bash
# Check IOptimizationPass interface definition
rg -nP --type=cs 'interface IOptimizationPass' -A 15

Length of output: 2529

---

🏁 Script executed:

#!/bin/bash
# Look for JitCompiler usage in compilation workflows
rg -nP --type=cs 'Compile|JitCompiler' src/JitCompiler/JitCompiler.cs -A 5 | head -100

Length of output: 4319

---

🏁 Script executed:

#!/bin/bash
# Search for any threading documentation in JitCompiler
rg -nP --type=cs 'Thread|thread.safe|concurrent' src/JitCompiler/JitCompiler.cs -C 3

Length of output: 0

---

Replace Dictionary with ConcurrentDictionary in AutoTuningPass._tuningCache.

The caching dictionary in AutoTuningPass is shared mutable state. Since JitCompiler creates a single instance of each optimization pass and reuses it across all compilations (line 135 & loop at line 451), concurrent calls to Compile() will invoke Optimize() on the same AutoTuningPass instance, causing race conditions on the unprotected dictionary. Replace with ConcurrentDictionary<int, TuningConfig> to match the thread-safety pattern already established in JitCompiler._compiledGraphCache.

🤖 Prompt for AI Agents

In src/JitCompiler/Optimizations/AutoTuningPass.cs around lines 81 to 103, the
_tuningCache is a plain Dictionary which is not thread-safe and can cause races
when Optimize() is called concurrently; replace the field type with
System.Collections.Concurrent.ConcurrentDictionary<int, TuningConfig>, update
the initializer accordingly, and change cache usage to use thread-safe methods
(e.g., TryGetValue plus GetOrAdd or AddOrUpdate) when reading/writing entries so
no locking is required.

[coderabbitai]

⚠️ Potential issue | 🟠 Major

Clarify or enforce the “no mutation of input graph” contract.

The XML docs state that Optimize must not modify the input IRGraph and should return a new, semantically equivalent graph. However, OperationFusionPass (snippet provided) appears to:

• Create a new IRGraph, but
• Mutate IROp instances (e.g., reassigning op.InputIds) that originated from graph.Operations.

That means callers holding on to the original IRGraph may observe mutated ops, which contradicts the documented behavior.

I’d recommend deciding on one of these and adjusting accordingly:

• Option A (pure contract): Keep the “no mutation” guarantee and update passes like OperationFusionPass to deep-copy ops (and any other mutable structures) before transforming them.
• Option B (imperative contract): Allow mutation of the input graph, but update the XML docs to say passes may mutate graph in-place and should document any side-effects.

Right now the interface documentation and at least one implementation are out of sync, which can be surprising for anyone implementing or composing passes.

[coderabbitai]

⚠️ Potential issue | 🟠 Major

Consider treating graph outputs as “uses” when deciding if an op can be safely fused.

All fusion patterns rely on CountUsages to require a single consumer, but CountUsages only counts operation inputs and ignores graph.OutputIds. That means an op whose tensor is both:

• consumed by a downstream op, and
• exposed as a graph output

still appears to have a single “use” and becomes eligible for fusion. After fusion, the tensor may no longer be produced as a distinct value, but OutputIds still reference it, which can break externally visible graph outputs.

To keep semantics safe you could:

• treat OutputIds as additional usages inside CountUsages, or
• explicitly prevent fusion if the candidate tensorId is in graph.OutputIds.

This is especially relevant for MatMul/Add/Activation and residual patterns where intermediate tensors may intentionally be marked as outputs.

Also applies to: 226-287, 351-393, 431-485, 487-499

[coderabbitai]

⚠️ Potential issue | 🔴 Critical

Bug: fused elementwise+activation op is being dropped from the optimized graph.

In FuseElementwiseActivation, you replace operations[i] with a new FusedElementwiseActivationOp and then do fusedOps.Add(operations[i]). At that point operations[i] is the fused op, so later in Optimize the fused op is treated as fused and skipped when building optimizedGraph.Operations. The original elementwise op is never added to fusedOps, so its output tensor is remapped but has no producer in the final graph.

This will remove the elementwise+activation computation entirely and leave downstream uses pointing to a tensor with no defining op.

A minimal fix is to mark the original elementwise op as fused, not the newly created one, e.g.:

-    private int FuseElementwiseActivation(List<IROp> operations, HashSet<IROp> fusedOps, Dictionary<int, int> tensorMapping)
+    private int FuseElementwiseActivation(List<IROp> operations, HashSet<IROp> fusedOps, Dictionary<int, int> tensorMapping)
     {
         int count = 0;

         for (int i = 0; i < operations.Count - 1; i++)
         {
             if (fusedOps.Contains(operations[i])) continue;

-            string? elementwiseOp = operations[i] switch
+            var originalElemwise = operations[i];
+            string? elementwiseOp = originalElemwise switch
             {
                 AddOp => "Add",
                 SubtractOp => "Subtract",
                 ElementwiseMultiplyOp => "Multiply",
                 DivideOp => "Divide",
                 _ => null
             };

             if (elementwiseOp == null) continue;
-            if (operations[i].InputIds.Length != 2) continue;
+            if (originalElemwise.InputIds.Length != 2) continue;

-            var elemwiseOutput = operations[i].OutputId;
+            var elemwiseOutput = originalElemwise.OutputId;
             ...
-                var fusedOp = new FusedElementwiseActivationOp
+                var fusedOp = new FusedElementwiseActivationOp
                 {
                     OutputId = operations[j].OutputId,
-                    InputIds = operations[i].InputIds,
+                    InputIds = originalElemwise.InputIds,
                     ...
                 };

-                operations[i] = fusedOp;
-                fusedOps.Add(operations[i]);
+                operations[i] = fusedOp;
+                fusedOps.Add(originalElemwise);
                 fusedOps.Add(operations[j]);

Any equivalent change that keeps the fused op out of fusedOps (so it’s emitted) and marks the original op as fused would work.

📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change

	private int FuseElementwiseActivation(List<IROp> operations, HashSet<IROp> fusedOps, Dictionary<int, int> tensorMapping)
	{
	int count = 0;
	for (int i = 0; i < operations.Count - 1; i++)
	{
	if (fusedOps.Contains(operations[i])) continue;
	string? elementwiseOp = operations[i] switch
	{
	AddOp => "Add",
	SubtractOp => "Subtract",
	ElementwiseMultiplyOp => "Multiply",
	DivideOp => "Divide",
	_ => null
	};
	if (elementwiseOp == null) continue;
	if (operations[i].InputIds.Length != 2) continue;
	var elemwiseOutput = operations[i].OutputId;
	// Find activation
	for (int j = i + 1; j < operations.Count; j++)
	{
	if (fusedOps.Contains(operations[j])) continue;
	string? activationName = operations[j] switch
	{
	ReLUOp => "ReLU",
	SigmoidOp => "Sigmoid",
	TanhOp => "Tanh",
	_ => null
	};
	if (activationName == null) continue;
	if (operations[j].InputIds.Length != 1 || operations[j].InputIds[0] != elemwiseOutput) continue;
	if (CountUsages(operations, elemwiseOutput, fusedOps) != 1) continue;
	// Create fused operation
	var fusedOp = new FusedElementwiseActivationOp
	{
	OutputId = operations[j].OutputId,
	InputIds = operations[i].InputIds,
	OutputType = operations[j].OutputType,
	OutputShape = operations[j].OutputShape,
	ElementwiseOp = elementwiseOp,
	ActivationName = activationName
	};
	operations[i] = fusedOp;
	fusedOps.Add(operations[i]);
	fusedOps.Add(operations[j]);
	tensorMapping[elemwiseOutput] = operations[j].OutputId;
	count++;
	break;
	}
	}
	private int FuseElementwiseActivation(List<IROp> operations, HashSet<IROp> fusedOps, Dictionary<int, int> tensorMapping)
	{
	int count = 0;
	for (int i = 0; i < operations.Count - 1; i++)
	{
	if (fusedOps.Contains(operations[i])) continue;
	var originalElemwise = operations[i];
	string? elementwiseOp = originalElemwise switch
	{
	AddOp => "Add",
	SubtractOp => "Subtract",
	ElementwiseMultiplyOp => "Multiply",
	DivideOp => "Divide",
	_ => null
	};
	if (elementwiseOp == null) continue;
	if (originalElemwise.InputIds.Length != 2) continue;
	var elemwiseOutput = originalElemwise.OutputId;
	// Find activation
	for (int j = i + 1; j < operations.Count; j++)
	{
	if (fusedOps.Contains(operations[j])) continue;
	string? activationName = operations[j] switch
	{
	ReLUOp => "ReLU",
	SigmoidOp => "Sigmoid",
	TanhOp => "Tanh",
	_ => null
	};
	if (activationName == null) continue;
	if (operations[j].InputIds.Length != 1 || operations[j].InputIds[0] != elemwiseOutput) continue;
	if (CountUsages(operations, elemwiseOutput, fusedOps) != 1) continue;
	// Create fused operation
	var fusedOp = new FusedElementwiseActivationOp
	{
	OutputId = operations[j].OutputId,
	InputIds = originalElemwise.InputIds,
	OutputType = operations[j].OutputType,
	OutputShape = operations[j].OutputShape,
	ElementwiseOp = elementwiseOp,
	ActivationName = activationName
	};
	operations[i] = fusedOp;
	fusedOps.Add(originalElemwise);
	fusedOps.Add(operations[j]);
	tensorMapping[elemwiseOutput] = operations[j].OutputId;
	count++;
	break;
	}
	}

[coderabbitai]

⚠️ Potential issue | 🟠 Major

🧩 Analysis chain

Fix boolean casing expectation in CompilationStats string test

CompilationStats.ToString() (per the snippet in JitCompiler.cs) formats the cache flag using the default bool formatter, which produces "True"/"False", whereas this test asserts Assert.Contains("false", str); (lowercase). That mismatch will cause the test to fail even when the implementation is correct.

Either:

• Update the test to match current behavior, e.g. Assert.Contains("False", str);, or
• Change CompilationStats.ToString() to emit a lowercase value (CacheHit.ToString().ToLowerInvariant()) and keep the test as‑is.

Given this is purely a presentation concern, aligning the test with the existing implementation is the least intrusive change.

---

I'll verify the boolean casing claim by examining the CompilationStats.ToString() implementation.

---

🏁 Script executed:

```shell
#!/bin/bash
# Find CompilationStats class and its ToString implementation
fd -type f -name "*.cs" | xargs rg -l "class CompilationStats" -A 50 | head -5

Length of output: 201

---

🏁 Script executed:

#!/bin/bash
# Search for CompilationStats class and ToString method with context
rg -n "class CompilationStats" -A 100 --type cs

Length of output: 5939

---

🏁 Script executed:

#!/bin/bash
# Verify the test file content matches the review snippet
rg -n "CompilationStats_ToString_ContainsRelevantInfo" -A 25 tests/AiDotNet.Tests/UnitTests/JitCompiler/JitCompilerTests.cs

Length of output: 901

---

Update test assertion to match CompilationStats boolean formatting

The test at line 268 asserts Assert.Contains("false", str); but CompilationStats.ToString() uses default C# bool formatting (line 653: $" Cache hit: {CacheHit}"), which produces capitalized "False", not "false". The test will fail.

Change line 268 to:

Assert.Contains("False", str);

🤖 Prompt for AI Agents

In tests/AiDotNet.Tests/UnitTests/JitCompiler/JitCompilerTests.cs around lines
248 to 269, the test asserts the string "false" but CompilationStats.ToString()
formats the boolean using C# default capitalization ("False"), causing the
assertion to fail; update the assertion to expect "False" (i.e., change the
Assert.Contains("false", str) to Assert.Contains("False", str")) so the test
matches the produced string.

[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)

2227-2245: ComputeGradients returns parameters instead of gradients

Although not part of the new JIT region, this method looks incorrect:

Backpropagate(outputGradients);

var gradients = new List<T>();
foreach (var layer in Layers)
{
    var layerParams = layer.GetParameters();
    gradients.AddRange(layerParams.ToArray());
}

After calling Backpropagate, it ignores parameter gradients and simply concatenates layer.GetParameters() (the weights) into the result. That makes ComputeGradients effectively return parameters, not gradients.

Given you already have GetParameterGradients() and GetGradients() helpers that use layer.GetParameterGradients(), this is very likely a bug.

Consider either:

• Replacing the loop with return GetParameterGradients();, or
• Calling layer.GetParameterGradients() instead of GetParameters().

Example minimal fix:

-    var gradients = new List<T>();
-    foreach (var layer in Layers)
-    {
-        var layerParams = layer.GetParameters();
-        gradients.AddRange(layerParams.ToArray());
-    }
-
-    return new Vector<T>(gradients.ToArray());
+    return GetParameterGradients();

🧹 Nitpick comments (6)

src/Models/VectorModel.cs (1)

1672-1760: VectorModel JIT export is structurally sound; add a simple training/coefficients guard

The JIT graph construction (input [1, FeatureCount] × coefficients [FeatureCount, 1]) matches the linear model and is consistent with the rest of the API. To avoid hard-to-debug failures when Coefficients is empty or uninitialized, consider mirroring the validation used in regression/time-series:

public ComputationNode<T> ExportComputationGraph(List<ComputationNode<T>> inputNodes)
{
-    if (inputNodes == null)
-        throw new ArgumentNullException(nameof(inputNodes));
+    if (inputNodes == null)
+        throw new ArgumentNullException(nameof(inputNodes));
+
+    if (Coefficients == null || Coefficients.Length == 0)
+    {
+        throw new InvalidOperationException(
+            "Cannot export computation graph: model has no coefficients (is it trained?).");
+    }

Optionally, you could also make SupportsJitCompilation reflect this state (e.g., return Coefficients != null && Coefficients.Length > 0) for consistency with TimeSeriesModelBase<T>.SupportsJitCompilation.

src/Models/NeuralNetworkModel.cs (1)

1245-1325: Verify tensor shapes/orientation for layer conversions and consider delegating to NeuralNetworkBase

The overall graph-construction flow is solid, but a few details are worth tightening up:

1. Dense layer weights/bias shapes

   ConvertDenseLayer uses:

   var weights = layer.GetWeights();
   var biases = layer.GetBiases();
   var weightsTensor = MatrixToTensor(weights);   // shape [rows, cols]
   var biasesTensor = VectorToTensor(biases);     // shape [biases.Length]
   var matmulNode = TensorOperations.MatrixMultiply(input, weightsNode);
   var addNode = TensorOperations.Add(matmulNode, biasesNode);

   While this may be correct, NeuralNetworkBase<T>.ConvertDenseLayer explicitly constructs weights as [inputSize, outputSize] and biases as [1, outputSize]. It would be good to confirm that:

   • DenseLayer<T>.GetWeights() already returns a matrix shaped for input @ weights (no transpose required), and
   • Tensor<T>.Add correctly broadcasts between [1, outputSize] and [outputSize] (if not, bias tensors should probably be shaped [1, outputSize] explicitly, as in the base implementation).

2. Conv/Pooling parameter fidelity

   ConvertConvolutionalLayer, ConvertMaxPoolingLayer, and ConvertAvgPoolingLayer currently hard-code or partially infer spatial hyperparameters:

   var stride = new int[] { 1, 1 };   // TODO for conv
   var padding = new int[] { 0, 0 };  // conv and pooling

   If your ConvolutionalLayer<T> / pooling layers support configurable stride or padding, JIT inference will diverge from normal inference. Plumb actual stride/padding from the layer APIs as soon as they’re available.

3. Avoiding duplicated JIT conversion logic

   NeuralNetworkBase<T> already exposes ExportComputationGraph(List<ComputationNode<T>>) and its own ConvertLayerToGraph, with similar responsibilities. To reduce drift between model-level and network-level JIT support, consider delegating:

   public ComputationNode<T> ExportComputationGraph(List<ComputationNode<T>> inputNodes)
   {
       if (inputNodes == null)
           throw new ArgumentNullException(nameof(inputNodes));
   
       return Network.ExportComputationGraph(inputNodes);
   }

   and, if possible, consolidating the layer-conversion helpers in one place instead of maintaining two parallel implementations.

Also applies to: 1341-1412, 1460-1507

docs/JIT_IMPLEMENTATION_STATUS.md (1)

15-25: Clarify why polynomial and Laplacian kernels are “unsupported”

The notes say:

• Polynomial ✗ (requires Power operation)
• Laplacian ✗ (requires Abs operation)

From the Autodiff side you already have generic Power/Abs-style primitives available; what’s actually missing is a JIT export path wiring those ops into NonLinearRegressionBase.ExportComputationGraph.

Consider rephrasing to something like “not yet wired into JIT export (would use Power/Abs ops)” so it’s clear the limitation is in the current graph export, not the underlying Autodiff capability.

src/NeuralNetworks/NeuralNetworkBase.cs (3)

2324-2411: SupportsJitCompilation should reflect actual layer support rather than always returning true

Right now:

public virtual bool SupportsJitCompilation => true;

but ConvertLayerToGraph only handles a subset of layer types and throws NotSupportedException for everything else. That means callers can see SupportsJitCompilation == true and still hit runtime failures when exporting a graph.

Consider having SupportsJitCompilation inspect the current Layers collection and return false when any layer is outside the supported set (Dense/FullyConnected/FeedForward/Activation/Dropout/GaussianNoise/Flatten/Reshape/Input/Masking/PositionalEncoding/BatchNorm/LayerNorm), or at least expose a helper that performs this check so orchestrating code can avoid invoking JIT on unsupported networks.

---

2420-2576: Reflection-heavy layer converters are fragile; add defensive checks or expose explicit JIT accessors

The converters for:

• FullyConnectedLayer<T> (_weights, _biases via GetField)
• FeedForwardLayer<T> (Weights, Biases via GetProperty)
• (and later) BatchNorm/LayerNorm private fields

all use reflection with null-forgiving operators:

var weightsField = layerType.GetField("_weights", BindingFlags.NonPublic | BindingFlags.Instance);
...
var weights = (Matrix<T>)weightsField!.GetValue(layer)!;

If the internal field/property names or visibilities change, this will surface as NullReferenceException inside JIT export, which is hard to diagnose from the outside.

At minimum, it would be safer to:

• Check for null and throw a clear NotSupportedException when the expected members are not found.
• Prefer strongly-typed, explicit APIs on the layer classes (e.g., internal/protected getters for weights/biases) over reflection when you control both sides.

Example pattern:

var weightsField = layerType.GetField("_weights", BindingFlags.NonPublic | BindingFlags.Instance);
if (weightsField == null)
{
    throw new NotSupportedException(
        $"JIT export for {layerType.Name} requires private field '_weights', which was not found. " +
        "This layer implementation may have changed; update the JIT converter accordingly.");
}

---

2581-2595: Activation mapping by type name is brittle and may miss custom/internal activations

ConvertActivationLayer dispatches based on layer.ActivationFunction.GetType().Name string values:

var activationType = layer.ActivationFunction.GetType().Name;

return activationType switch
{
    "ReLU" or "ReLUActivation" => TensorOperations.ReLU(input),
    "Sigmoid" or "SigmoidActivation" => TensorOperations.Sigmoid(input),
    ...

This is sensitive to exact type names and makes it hard to support:

• Custom activation implementations.
• Renamed/broken-up builtin activations.
• Multiple aliases beyond the few hardcoded strings.

If possible, prefer:

• An enum or capability-centric API on the activation (e.g., ActivationKind property), or
• An interface with well-known types you can match on directly (is ReLUActivation<T>, etc.)

Even a small wrapper method on ActivationLayer<T> that exposes an ActivationFunctionType enum would make this mapping much more robust.

📜 Review details

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 2371f17 and 124dfbe.

📒 Files selected for processing (8)

• docs/JIT_IMPLEMENTATION_STATUS.md (1 hunks)
• src/Interfaces/IFullModel.cs (1 hunks)
• src/Models/NeuralNetworkModel.cs (3 hunks)
• src/Models/VectorModel.cs (2 hunks)
• src/NeuralNetworks/NeuralNetworkBase.cs (2 hunks)
• src/Regression/NonLinearRegressionBase.cs (2 hunks)
• src/Regression/RegressionBase.cs (2 hunks)
• src/TimeSeries/TimeSeriesModelBase.cs (2 hunks)

🧰 Additional context used

🧬 Code graph analysis (7)

src/Interfaces/IFullModel.cs (1)

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

• TOutput (626-663)

src/Regression/RegressionBase.cs (2)

src/Models/VectorModel.cs (3)

• ComputationNode (1723-1745)
• T (382-401)
• Tensor (1750-1760)

src/Autodiff/TensorOperations.cs (1)

• TensorOperations (45-5389)

src/TimeSeries/TimeSeriesModelBase.cs (3)

src/Regression/RegressionBase.cs (7)

• ComputationNode (1001-1046)
• Vector (168-178)
• Vector (247-250)
• Vector (347-360)
• Vector (380-388)
• Vector (410-429)
• Vector (793-829)

src/Models/VectorModel.cs (7)

• ComputationNode (1723-1745)
• T (382-401)
• Tensor (1750-1760)
• Vector (278-312)
• Vector (521-540)
• Vector (890-898)
• Vector (922-932)

src/Autodiff/TensorOperations.cs (1)

• TensorOperations (45-5389)

src/Regression/NonLinearRegressionBase.cs (4)

src/NeuralNetworks/NeuralNetworkBase.cs (22)

• ComputationNode (2377-2411)
• ComputationNode (2420-2448)
• ComputationNode (2453-2494)
• ComputationNode (2499-2542)
• ComputationNode (2547-2576)
• ComputationNode (2581-2596)
• ComputationNode (2601-2607)
• ComputationNode (2612-2676)
• ComputationNode (2681-2715)
• T (1056-1065)
• Tensor (250-271)
• Tensor (287-290)
• Tensor (366-386)
• Tensor (394-401)
• Tensor (421-443)
• Tensor (480-525)
• Tensor (873-873)
• Tensor (1163-1185)
• Vector (306-309)
• Vector (321-343)
• Vector (817-837)
• Vector (2134-2151)

src/TimeSeries/TimeSeriesModelBase.cs (4)

• ComputationNode (1781-1823)
• T (428-428)
• T (1068-1111)
• T (1342-1355)

src/Regression/RegressionBase.cs (7)

• ComputationNode (1001-1046)
• Vector (168-178)
• Vector (247-250)
• Vector (347-360)
• Vector (380-388)
• Vector (410-429)
• Vector (793-829)

src/Autodiff/TensorOperations.cs (1)

• TensorOperations (45-5389)

src/Models/NeuralNetworkModel.cs (2)

src/NeuralNetworks/NeuralNetworkBase.cs (17)

• ComputationNode (2377-2411)
• ComputationNode (2420-2448)
• ComputationNode (2453-2494)
• ComputationNode (2499-2542)
• ComputationNode (2547-2576)
• ComputationNode (2581-2596)
• ComputationNode (2601-2607)
• ComputationNode (2612-2676)
• T (1056-1065)
• Tensor (250-271)
• Tensor (287-290)
• Tensor (366-386)
• Tensor (394-401)
• Tensor (421-443)
• Tensor (480-525)
• Tensor (873-873)
• Tensor (1163-1185)

src/Interfaces/IJitCompilable.cs (1)

• ComputationNode (84-84)

src/Models/VectorModel.cs (4)

src/TimeSeries/TimeSeriesModelBase.cs (5)

• ComputationNode (1781-1823)
• T (428-428)
• T (1068-1111)
• T (1342-1355)
• Vector (358-379)

src/Regression/RegressionBase.cs (1)

• ComputationNode (1001-1046)

src/Models/NeuralNetworkModel.cs (13)

• ComputationNode (1244-1262)
• ComputationNode (1267-1292)
• ComputationNode (1294-1325)
• ComputationNode (1327-1339)
• ComputationNode (1341-1372)
• ComputationNode (1374-1382)
• ComputationNode (1384-1392)
• ComputationNode (1394-1412)
• ComputationNode (1414-1426)
• ComputationNode (1428-1436)
• Tensor (546-553)
• Tensor (1489-1493)
• Tensor (1498-1507)

src/Autodiff/TensorOperations.cs (1)

• TensorOperations (45-5389)

src/NeuralNetworks/NeuralNetworkBase.cs (3)

src/Regression/RegressionBase.cs (7)

• ComputationNode (1001-1046)
• Vector (168-178)
• Vector (247-250)
• Vector (347-360)
• Vector (380-388)
• Vector (410-429)
• Vector (793-829)

src/Autodiff/ComputationNode.cs (3)

• ComputationNode (28-414)
• ComputationNode (212-225)
• List (301-342)

src/Autodiff/TensorOperations.cs (11)

• ComputationNode (70-81)
• ComputationNode (101-104)
• ComputationNode (127-176)
• ComputationNode (203-251)
• ComputationNode (278-326)
• ComputationNode (353-418)
• ComputationNode (443-484)
• ComputationNode (505-539)
• ComputationNode (559-597)
• List (3201-3279)
• TensorOperations (45-5389)

🔇 Additional comments (1)

src/Interfaces/IFullModel.cs (1)

43-45: IJitCompilable inheritance is a breaking change for all IFullModel implementations

Making IFullModel<T, TInput, TOutput> inherit IJitCompilable<T, TInput, TOutput> means every existing implementation (including external ones) must now provide SupportsJitCompilation and ExportComputationGraph(...). Please double-check that:

• All in-repo implementations compile and have correct JIT graph behavior.
• This interface expansion is called out in release notes for consumers who implement IFullModel<, ,> themselves.

[coderabbitai]

⚠️ Potential issue | 🟠 Major

JIT docs vs behavior: AddLayer/ConcatenateLayer are declared supported but treated as no-ops

The JIT documentation block lists AddLayer and ConcatenateLayer as supported, and ConvertLayerToGraph explicitly routes these types to ConvertAddLayer / ConvertConcatenateLayer. However both implementations currently just return the input node:

private ComputationNode<T> ConvertAddLayer(AddLayer<T> layer, ComputationNode<T> input)
{
    // ... comment about residual connections ...
    return input;
}

private ComputationNode<T> ConvertConcatenateLayer(ConcatenateLayer<T> layer, ComputationNode<T> input)
{
    // ... comment about multiple inputs ...
    return input;
}

For architectures that actually use Add/Concat (residual blocks, skip connections, multi-branch subgraphs), this will silently change the network semantics in the JIT path and produce incorrect predictions.

Until proper multi-input graph handling is wired up, it would be safer to fail fast rather than approximate:

private ComputationNode<T> ConvertAddLayer(AddLayer<T> layer, ComputationNode<T> input)
{
-    // AddLayer typically adds a residual connection
-    // This requires multiple inputs which isn't supported in simple forward pass
-    // For now, just return input (residual connections need graph restructuring)
-    return input;
+    throw new NotSupportedException(
+        "AddLayer<T> is not yet supported for JIT compilation. " +
+        "Residual connections require multi-input graph wiring. " +
+        "Please disable JIT or remove AddLayer<T> layers for now.");
}

private ComputationNode<T> ConvertConcatenateLayer(ConcatenateLayer<T> layer, ComputationNode<T> input)
{
-    // Concatenation requires multiple inputs
-    // For simple forward pass, just return input
-    // Full support requires restructuring the graph to handle multiple inputs
-    return input;
+    throw new NotSupportedException(
+        "ConcatenateLayer<T> is not yet supported for JIT compilation. " +
+        "Concatenation requires multi-input graph wiring. " +
+        "Please disable JIT or avoid ConcatenateLayer<T> layers for now.");
}

If you prefer to keep them in the supported set, the docs should be updated to clarify current limitations (e.g., “not yet JIT-supported”) so users don’t infer full fidelity.

Also applies to: 1183-1183, 1444-1458

🤖 Prompt for AI Agents

In src/Models/NeuralNetworkModel.cs around lines 17 to 46 (and also affecting
ConvertAddLayer/ConvertConcatenateLayer at ~1183 and ~1444-1458), the JIT docs
claim AddLayer/ConcatenateLayer are supported but their converters are no-ops
that return the input node, silently changing semantics for models that use
residuals or multi-branch concatenation; either (A) change the converters to
throw a clear NotSupportedException (or InvalidOperationException) when
encountering AddLayer/ConcatenateLayer so the JIT path fails fast, or (B) update
the documentation/comment block to remove them from the supported list (or mark
them “not yet JIT-supported”) and add TODO comments pointing to the multi-input
graph work required; implement one of these fixes consistently in all three
locations.

[coderabbitai]

⚠️ Potential issue | 🔴 Critical

Kernel computation graph is missing reduction, breaking semantic equivalence and risking shape errors

ComputeLinearKernel, ComputeRBFKernel, and ComputeSigmoidKernel currently return elementwise tensors (ElementwiseMultiply / squared differences) without any reduction over the feature dimension, but they are used as scalar kernels in:

• weightedNode = ElementwiseMultiply(kernelNode, alphaNode) where alphaNode is [1,1]
• Accumulation into sumNode
• Final Add(sumNode!, biasNode) where biasNode is [1,1]

Impact:

• If TensorOperations.Add/ElementwiseMultiply do not support broadcasting between [1, featureCount] and [1,1], this will throw at runtime in the JIT path.
• Even if broadcasting is supported, the JIT-exported model computes a vector-valued “kernel” instead of the scalar kernel used by PredictSingle/KernelFunction, so JIT predictions will not match non-JIT predictions.

These methods should explicitly reduce across the feature dimension to produce a scalar-like node (e.g., shape [1,1]) before weighting and summation.

A concrete fix could look like this (illustrative):

-    private ComputationNode<T> ComputeLinearKernel(ComputationNode<T> x1, ComputationNode<T> x2)
-    {
-        // Element-wise multiply
-        var product = TensorOperations.ElementwiseMultiply(x1, x2);
-
-        // Sum all elements (reduction)
-        // Note: For now, we'll use a simple approach
-        // In a full implementation, we'd have a proper Sum/Reduce operation
-        return product; // Simplified - assumes proper reduction in code generation
-    }
+    private ComputationNode<T> ComputeLinearKernel(ComputationNode<T> x1, ComputationNode<T> x2)
+    {
+        // Element-wise multiply
+        var product = TensorOperations.ElementwiseMultiply(x1, x2);
+
+        // Reduce across feature dimension to obtain scalar kernel value.
+        // Shape: [1, 1] for a single example.
+        var dot = TensorOperations.Sum(product, new[] { 1 }, keepDims: true);
+        return dot;
+    }
@@
-        // Sum squared differences (||x1 - x2||^2)
-        // Simplified - assumes proper reduction
-        var sumSquared = squared;
+        // Sum squared differences (||x1 - x2||^2) across features
+        var sumSquared = TensorOperations.Sum(squared, new[] { 1 }, keepDims: true);
@@
-        // Dot product: x1 · x2
-        var dotProduct = TensorOperations.ElementwiseMultiply(x1, x2);
-        // Simplified - assumes proper reduction
+        // Dot product: x1 · x2
+        var product = TensorOperations.ElementwiseMultiply(x1, x2);
+        // Reduce across feature dimension to obtain scalar dot product
+        var dotProduct = TensorOperations.Sum(product, new[] { 1 }, keepDims: true);

This keeps the exported graph consistent with the scalar kernel semantics used elsewhere and avoids relying on undocumented broadcasting behavior in TensorOperations.

Committable suggestion skipped: line range outside the PR's diff.

🤖 Prompt for AI Agents

In src/Regression/NonLinearRegressionBase.cs around lines 1200 to 1338, the
kernel helper methods return elementwise feature-length tensors instead of
scalar kernel values, causing shape/broadcasting and semantic mismatches; update
ComputeLinearKernel to reduce the elementwise product across the feature
dimension (sum/reduce to shape [1,1]) before returning, update ComputeRBFKernel
to sum the squared differences to a scalar then multiply by -gamma and exp, and
update ComputeSigmoidKernel to reduce the dot-product to a scalar before
scaling, adding coef0 and applying tanh; use the project's TensorOperations
reduction API (e.g., Sum/Reduce with axis for features) so all kernel nodes are
shape [1,1] prior to multiplication by alpha and accumulation.

[coderabbitai]

⚠️ Potential issue | 🔴 Critical

Guard against null inputNodes in RegressionBase ExportComputationGraph

The JIT graph export for regression looks correct (X @ w plus optional intercept), but inputNodes is assumed non-null and used directly:

var inputNode = new ComputationNode<T>(inputTensor);
inputNodes.Add(inputNode);

To avoid a NullReferenceException and match the pattern used elsewhere (e.g., VectorModel<T>):

public virtual ComputationNode<T> ExportComputationGraph(List<ComputationNode<T>> inputNodes)
{
-    // Validation: Ensure model is trained
+    if (inputNodes == null)
+    {
+        throw new ArgumentNullException(nameof(inputNodes));
+    }
+
+    // Validation: Ensure model is trained
     if (Coefficients == null || Coefficients.Length == 0)
     {
         throw new InvalidOperationException("Cannot export computation graph: Model has not been trained yet.");
     }

🤖 Prompt for AI Agents

In src/Regression/RegressionBase.cs around lines 952 to 1045,
ExportComputationGraph assumes inputNodes is non-null and calls
inputNodes.Add(...) which can throw NullReferenceException; add an explicit
guard at the start of the method that checks if inputNodes is null and throws
ArgumentNullException(nameof(inputNodes)) (matching the pattern used in
VectorModel<T>), so the method fails fast with a clear error rather than causing
a null reference later.

[coderabbitai]

⚠️ Potential issue | 🔴 Critical

Add a null-check for inputNodes in ExportComputationGraph to avoid NRE

ExportComputationGraph assumes inputNodes is non-null and calls inputNodes.Add(inputNode) unconditionally. If a caller accidentally passes null, this will surface as a NullReferenceException rather than a clear argument error.

Recommend adding a guard:

public virtual ComputationNode<T> ExportComputationGraph(List<ComputationNode<T>> inputNodes)
{
-    // Validation: Ensure model is trained
+    if (inputNodes == null)
+    {
+        throw new ArgumentNullException(nameof(inputNodes));
+    }
+
+    // Validation: Ensure model is trained
     if (!IsTrained)
     {
         throw new InvalidOperationException("Cannot export computation graph: Model has not been trained yet.");
     }

🤖 Prompt for AI Agents

In src/TimeSeries/TimeSeriesModelBase.cs around lines 1719 to 1823,
ExportComputationGraph assumes inputNodes is non-null and will throw a
NullReferenceException when calling inputNodes.Add(...); add a guard at the
start of the method to validate inputNodes and throw an ArgumentNullException
(with the parameter name "inputNodes") if it is null, then proceed with the
existing logic; keep the check before any use of inputNodes so callers get a
clear, descriptive argument error instead of an NRE.