fix: address review comments for mixed-precision training

[ooples]

Summary

• LayerBase: fix Clone() to assign fresh instance ID for unique LayerName
• MixedPrecisionTrainingLoop: throw for None/TF32 precision types (not supported)
• MixedPrecisionScope: add RegisterFP32Only method for FP32-only layers
• LayerBase: use RegisterFP32Only to avoid unnecessary FP16 copies for FP32 layers
• GradientBasedOptimizerBase: cache gradients unscaled to handle dynamic loss scaling correctly
• MixedPrecisionIntegrationTests: update LayerName assertion to use StartsWith pattern

Test plan

• Build passes
• Existing mixed-precision tests pass
• LayerName tests updated for new instance ID suffix pattern

🤖 Generated with Claude Code

[vercel]

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

Project	Deployment	Actions	Updated (UTC)
aidotnet-playground-api	Ready	Preview, Comment	Feb 4, 2026 5:43pm

[coderabbitai]

Summary by CodeRabbit

Release Notes

• New Features
  • Added 8-bit floating-point support with two precision variants, enabling reduced memory footprint for model training and inference
  • Introduced conversion utilities for seamless interoperability between different precision formats
  • Enhanced gradient handling in optimizers to support unscaled gradient application for improved numerical stability in mixed-precision workflows

Walkthrough

The pull request introduces FP8 numeric types (Float8E4M3, Float8E5M2) with conversion utilities, updates layer precision policy logic to support FP8-type detection, and enhances gradient application in optimizers with unscaled gradient handling.

Changes

Cohort / File(s)	Summary
FP8 Type System src/MixedPrecision/Float8Types.cs	Adds two public 8-bit FP8 types with bit-level conversions: Float8E4M3 (4 exponent, 3 mantissa) and Float8E5M2 (5 exponent, 2 mantissa), each with value constants, conversion methods, properties (IsNaN, IsZero, IsNegative, IsInfinity for E5M2), and comparison operators. Includes internal BitConverterHelper and public Float8Extensions with array conversions and cross-format conversion methods.
Layer Precision Policy src/MixedPrecision/LayerPrecisionPolicy.cs	Broadens ShouldUseHigherPrecision logic to handle FP8 types by introducing private IsFp8() helper that categorizes FP8-related precisions. Logic now returns true for None (FP32) layers and checks if default precision is FP8 when layer is FP16 or BF16.
Gradient Application src/Optimizers/GradientBasedOptimizerBase.cs	Adds private GetUnscaledGradients() helper and new public ApplyGradients() overload accepting gradientsAlreadyUnscaled parameter to avoid double-unscaling. Adjusts gradient-statistics helpers (vanishing/exploding checks, norm calculation) to operate on unscaled gradients when mixed precision is active.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~25 minutes

Possibly related PRs

• Claude/mixed precision training architecture 011 cv13wm js ax6s gj6 ryu37k #479: Adds mixed-precision-related gradient handling methods to GradientBasedOptimizerBase.cs, directly overlapping with the gradient optimizer modifications in this PR.

Poem

🐰 Eight bits of magic, hopping through the code,
E4M3 and E5M2 lighten the load,
Unscaled gradients flow so true and clean,
The smallest precision we've ever seen! ✨

🚥 Pre-merge checks | ✅ 2 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 56.20% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (2 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title 'fix: address review comments for mixed-precision training' is directly related to the changeset, which addresses multiple review-comment fixes for mixed-precision training including Float8 types, LayerPrecisionPolicy, and GradientBasedOptimizerBase.
Description check	✅ Passed	The description details specific changes across multiple components (LayerBase, MixedPrecisionTrainingLoop, MixedPrecisionScope, GradientBasedOptimizerBase, MixedPrecisionIntegrationTests) that align with the changeset, providing meaningful context about the fixes and test plan.

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

✨ Finishing touches

• 📝 Generate docstrings

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Post copyable unit tests in a comment
• Commit unit tests in branch fix/mixed-precision-review-comments

---

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

[Copilot]

Pull request overview

This PR implements mixed-precision training enhancements for neural networks, adding support for FP8 precision types, layer-specific precision policies, and fixing the Clone() method to generate unique instance IDs. The changes span across optimizer, layer, and mixed-precision infrastructure.

Changes:

• Added FP8 type implementations (Float8E4M3, Float8E5M2) with conversion utilities
• Implemented LayerPrecisionPolicy for granular layer-level precision control
• Added RegisterFP32Only method to MixedPrecisionScope to optimize FP32-only layer handling
• Updated LayerBase.Clone() to assign fresh instance IDs for unique LayerName generation
• Modified optimizer to cache gradients unscaled for correct dynamic loss scaling
• Added validation to reject None/TF32 precision types in MixedPrecisionTrainingLoop

Reviewed changes

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

Show a summary per file

File	Description
MixedPrecisionIntegrationTests.cs	Comprehensive test coverage for FP8 types, layer precision policies, scope management, and AMP integration
ContinualLearningTestHelper.cs	Updated MockLayer to implement new ILayer interface members (ForwardWithPrecisionCheck, LayerName)
TrainingMemoryManager.cs	Updated to use ForwardWithPrecisionCheck for mixed-precision aware forward passes
GradientBasedOptimizerBase.cs	Added gradient scaling/unscaling logic and routing for mixed-precision, extracted ApplyGradientsCore method
NeuralNetworkBase.cs	Updated ForwardWithMemory to use ForwardWithPrecisionCheck for layer-level precision awareness
LayerBase.cs	Added instance counter for unique LayerName, implemented ForwardWithPrecisionCheck, updated Clone() to assign fresh IDs
MixedPrecisionTrainingLoop.cs	Added policy parameter, validation for unsupported precision types, improved gradient scaling workflow
MixedPrecisionScope.cs	New file implementing ambient context for mixed-precision operations with FP16/FP32 tensor tracking
MixedPrecisionConfig.cs	Added FP8-specific settings and factory methods (ForFP8, ForBF16, Conservative, etc.)
LayerPrecisionPolicy.cs	New file implementing layer-specific precision policies with pattern matching and factory methods
Float8Types.cs	New file implementing FP8 E4M3 and E5M2 types with conversion utilities
ILayer.cs	Extended interface with ForwardWithPrecisionCheck method and LayerName property
MixedPrecisionType.cs	Added FP8_E4M3, FP8_E5M2, and FP8_Hybrid enum values with documentation

---

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

[coderabbitai]

Actionable comments posted: 1

🤖 Fix all issues with AI agents

In `@src/MixedPrecision/LayerPrecisionPolicy.cs`:
- Around line 57-71: Add input validation to LayerPrecisionPolicy.AddPattern to
prevent empty or all-whitespace patterns which would match every layer; before
adding to the _patterns list, trim the pattern and if
string.IsNullOrWhiteSpace(pattern) throw an ArgumentException (or
ArgumentNullException for null) with a clear message indicating pattern must be
non-empty/non-whitespace, otherwise proceed to _patterns.Add((pattern,
precision)) and return this.

🧹 Nitpick comments (6)

tests/AiDotNet.Tests/Helpers/ContinualLearningTestHelper.cs (1)

356-363: Consider making LayerName unique per instance.

Right now both mock layers return the same name, which can mask policy-lookup issues in tests. A simple instance suffix keeps the mock aligned with the interface guidance.

♻️ Suggested tweak

 public class MockLayer<T> : ILayer<T>
 {
+    private static int _instanceCounter;
+    private readonly int _instanceId = ++_instanceCounter;
@@
-    public string LayerName => "MockLayer";
+    public string LayerName => $"{GetType().Name}_{_instanceId}";

src/MixedPrecision/MixedPrecisionConfig.cs (2)

22-177: Guard FP8 format fields against non‑FP8 values.
FP8ForwardFormat/FP8BackwardFormat accept any MixedPrecisionType, so a misconfigured value (e.g., BF16) would only fail later in the pipeline. Consider a lightweight validation step when the config is consumed (or a helper on the config) to fail fast when PrecisionType is FP8 but the formats are not.

---

308-315: Consider including FP8 fields in ToString() for diagnostics.
When FP8 is in play, having forward/backward formats and exclusions in the summary string makes debug logs much more informative.

src/MixedPrecision/MixedPrecisionScope.cs (2)

205-225: Clear any stale FP16 copy when registering FP32‑only tensors.
If a name was previously registered via RegisterAndCastToFP16, _fp16Tensors retains the old entry. Removing it here avoids unexpected GetFP16Tensor results when names are reused.

♻️ Suggested diff

 public void RegisterFP32Only(string name, Tensor<float> fp32Tensor)
 {
     if (fp32Tensor == null)
     {
         throw new ArgumentNullException(nameof(fp32Tensor));
     }
 
     // Store only the FP32 version - no FP16 copy needed for FP32-only layers
     _fp32Tensors[name] = fp32Tensor;
+    _fp16Tensors.Remove(name);
 }

---

331-346: Defensive guard on scope restoration.
If scopes are disposed out of order, _current could be clobbered. A quick ReferenceEquals(Current, this) check before restoring _previous makes the behavior safer.

♻️ Suggested diff

-        // Restore previous scope
-        SetCurrentScope(_previous);
+        // Restore previous scope only if this scope is current
+        if (ReferenceEquals(Current, this))
+        {
+            SetCurrentScope(_previous);
+        }

src/MixedPrecision/LayerPrecisionPolicy.cs (1)

317-326: Optional: add null guards for public extension inputs.
A fast ArgumentNullException for config/policy makes misuse fail predictably.

[coderabbitai]

Actionable comments posted: 4

🤖 Fix all issues with AI agents

In `@src/MixedPrecision/Float8Types.cs`:
- Around line 171-198: FromFloat currently collapses FP32 subnormals and can
round large values into the NaN encoding; fix by handling e4m3Exponent <= 0 as a
subnormal path: build the 24-bit significand (floatMantissa | 0x800000),
right-shift it by (1 - e4m3Exponent) to produce a denormalized mantissa, then
take the top 3 bits (with proper rounding) into e4m3Mantissa instead of directly
using floatMantissa bits; for overflow saturation ensure you never emit the NaN
pattern by setting e4m3Exponent = 15 and e4m3Mantissa = 6 (the maximum finite
mantissa) if overflow or if rounding carry would produce mantissa == 7, adjust
carry handling so it promotes exponent but then clamps mantissa to 6 to avoid
the NaN encoding (references: FromFloat, e4m3Exponent, e4m3Mantissa,
MantissaMask, ExponentBias).

In `@src/MixedPrecision/LayerPrecisionPolicy.cs`:
- Around line 128-133: The current ShouldUseHigherPrecision uses enum ordering
and ignores BF16; update it to use explicit enum comparisons: call
GetPrecision(layerName) and return true if it's MixedPrecisionType.None (FP32),
or if it's MixedPrecisionType.FP16 or MixedPrecisionType.BF16 and
_defaultPrecision == MixedPrecisionType.FP8_E4M3 (i.e., only use higher
precision for FP16/BF16 when default is FP8); alternatively, if BF16 is
unsupported, explicitly treat MixedPrecisionType.BF16 as a blocker (return false
or throw) instead of relying on ordinal comparisons. Ensure you reference
ShouldUseHigherPrecision, GetPrecision, MixedPrecisionType, and
_defaultPrecision when making the change.

In `@src/Optimizers/GradientBasedOptimizerBase.cs`:
- Around line 329-343: Add a defensive guard so the optimizer does not
double-unscale gradients: introduce and check a flag/property (e.g.,
_gradientsAlreadyUnscaled or a public parameter like bool
gradientsAlreadyUnscaled) in the ApplyGradients entry point and in
ApplyGradientsWithMixedPrecision; if gradientsAlreadyUnscaled is true, skip
unscale/unmix logic in ApplyGradientsWithMixedPrecision and route directly to
ApplyGradientsCore, and ensure callers like MixedPrecisionTrainingLoop set that
flag when they perform unscaling so the optimizer (which uses
_mixedPrecisionContext) will not unscale again.
- Around line 652-669: The health checks are running on SCALED gradients because
_lastComputedGradients is set after ApplyMixedPrecisionScaling, causing false
positives; update the checks in AreGradientsExploding and AreGradientsVanishing
to operate on unscaled gradients by dividing _lastComputedGradients by the
current loss scale (or calling a dedicated Unscale helper) before applying the
thresholds, or alternatively make those methods accept a flag/currentLossScale
to adjust thresholds when mixed-precision is active; reference
_lastComputedGradients, ApplyMixedPrecisionScaling, AreGradientsExploding,
AreGradientsVanishing and ApplyGradientsWithMixedPrecision when implementing the
unscale or threshold-adjustment logic so checks always use the unscaled gradient
magnitude.

[Copilot]

Pull request overview

Copilot reviewed 13 out of 13 changed files in this pull request and generated 4 comments.

---

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

[Copilot]

The static readonly fields MaxValue and MinPositive in Float8E4M3 and Float8E5M2 structs create a circular dependency during type initialization. These fields call FromFloat() which is a static method, but the type's static constructor hasn't necessarily finished executing yet. While this may work in practice due to C# static initialization ordering, it's considered an anti-pattern and can lead to subtle initialization issues.

Consider using lazy initialization or static properties with backing fields to avoid potential type initialization order problems.

[Copilot]

The MixedPrecisionTrainingLoop unscales gradients at line 195, then passes them to the optimizer at line 207. If the optimizer has mixed-precision enabled (via _mixedPrecisionContext), it will attempt to unscale the gradients again in ApplyGradientsWithMixedPrecision, leading to double-unscaling and incorrect parameter updates.

While the comment at lines 332-335 in GradientBasedOptimizerBase.cs mentions this issue, the training loop should use the gradientsAlreadyUnscaled parameter (from the new overload at line 355) to explicitly prevent double-unscaling. Consider calling:

_optimizer.ApplyGradients(parameters, gradients, _network, gradientsAlreadyUnscaled: true)

This ensures correctness even if the optimizer has mixed-precision enabled.

Suggested change

	var updatedModel = _optimizer.ApplyGradients(parameters, gradients, _network);
	var updatedModel = _optimizer.ApplyGradients(parameters, gradients, _network, gradientsAlreadyUnscaled: true);

[Copilot]

The static _instanceCounter field in LayerBase could potentially overflow after 2,147,483,647 layer instances are created, causing duplicate instance IDs. While this is an extremely unlikely scenario in practice, consider using a long instead of int for the counter, or document this limitation.

This is a very minor issue that would only occur in extremely long-running applications creating billions of layer instances.

Suggested change

	private static int _instanceCounter;
	/// <summary>
	/// The unique instance ID for this layer, used to distinguish multiple instances of the same layer type.
	/// </summary>
	private int _instanceId;
	private static long _instanceCounter;
	/// <summary>
	/// The unique instance ID for this layer, used to distinguish multiple instances of the same layer type.
	/// </summary>
	private long _instanceId;

[coderabbitai]

Actionable comments posted: 1

Caution

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

⚠️ Outside diff range comments (1)

src/Optimizers/GradientBasedOptimizerBase.cs (1)

1071-1121: ⚠️ Potential issue | 🟡 Minor

Inconsistent mixed-precision handling in batch-based gradient calculation.

The batch-based CalculateGradient overload (lines 1071-1121) stores gradients directly in _lastComputedGradients without applying ApplyMixedPrecisionScaling. This is inconsistent with the main CalculateGradient method (line 705) which scales gradients before storing.

When mixed-precision is enabled and this batch-based method is used, GetUnscaledGradients() will incorrectly divide unscaled gradients by the loss scale, producing wrong values for health checks.

🔧 Proposed fix to apply consistent scaling

         // Average the gradient using vectorized division
         var batchSizeScalar = NumOps.FromDouble(batchIndices.Length);
         gradient = (Vector<T>)Engine.Divide(gradient, batchSizeScalar);

+        // Apply gradient clipping if enabled
+        gradient = ApplyGradientClipping(gradient);
+
+        // Scale gradients for mixed-precision training
+        gradient = ApplyMixedPrecisionScaling(gradient);
+
         // Store for external access (enables gradient clipping, true DDP, debugging, etc.)
         _lastComputedGradients = gradient;

         return gradient;

🤖 Fix all issues with AI agents

In `@src/MixedPrecision/Float8Types.cs`:
- Around line 437-443: The E5M2 FromFloat implementation incorrectly flushes
subnormals to zero: adjust the FromFloat logic in the E5M2 type to compute
denormalized mantissa when e5m2Exponent < 0 using the same pattern as E4M3
(shift the float mantissa right by (1 - e5m2Exponent) and set exponent to 0
instead of just clamping to 0), and update the E5M2 minVal constant to the real
minimum representable subnormal (use 2^(1-ExponentBias-NumMantissaBits)
equivalent to the E5M2 layout) so encoding/decoding (FromFloat and ToFloat) are
symmetric. Ensure you modify the code paths in FromFloat that clamp e5m2Exponent
and the minVal check together so values between true subnormal range are encoded
as denormals rather than flushed to zero.

🧹 Nitpick comments (7)

src/MixedPrecision/LayerPrecisionPolicy.cs (3)

51-55: Add input validation for layerName to match AddPattern consistency.

AddPattern validates against null/whitespace patterns, but SetPrecision does not validate layerName. A null value would cause a NullReferenceException from the dictionary, and an empty string creates a potentially confusing entry.

🛠️ Suggested fix

 public LayerPrecisionPolicy SetPrecision(string layerName, MixedPrecisionType precision)
 {
+    if (string.IsNullOrWhiteSpace(layerName))
+    {
+        throw new ArgumentException("Layer name cannot be null or whitespace.", nameof(layerName));
+    }
+
     _exactMatches[layerName] = precision;
     return this;
 }

---

103-121: Guard against null layerName to prevent NullReferenceException.

If layerName is null, _exactMatches.TryGetValue and layerName.Contains will throw. Consider adding a guard clause or returning the default precision for null/empty layer names.

🛠️ Suggested fix

 public MixedPrecisionType GetPrecision(string layerName)
 {
+    if (string.IsNullOrEmpty(layerName))
+    {
+        return _defaultPrecision;
+    }
+
     // Check exact matches first
     if (_exactMatches.TryGetValue(layerName, out var precision))
     {

---

357-366: Add null guards for extension method parameters.

Extension methods can be called with null receivers in C#. Consider adding null checks for both config and policy to provide better error messages.

🛠️ Suggested fix

 public static MixedPrecisionConfig WithLayerPolicy(this MixedPrecisionConfig config, LayerPrecisionPolicy policy)
 {
+    ArgumentNullException.ThrowIfNull(config);
+    ArgumentNullException.ThrowIfNull(policy);
+
     // Get excluded patterns from the policy based on the config's precision type
     var excludedPatterns = policy.GetExcludedPatterns(config.PrecisionType);

src/Optimizers/GradientBasedOptimizerBase.cs (1)

725-735: Consider adding defensive validation for the scale value.

GetUnscaledGradients() (lines 784-787) validates the scale against invalid values (<= 0, NaN, Infinity), but ApplyMixedPrecisionScaling does not. For consistency and defensive programming, consider adding the same validation here:

♻️ Proposed fix for consistent validation

 protected virtual Vector<T> ApplyMixedPrecisionScaling(Vector<T> gradient)
 {
     if (_mixedPrecisionContext == null)
     {
         return gradient;
     }

     // Scale gradient by the loss scale factor using efficient vector multiplication
     double scale = _mixedPrecisionContext.LossScaler.Scale;
+    if (scale <= 0 || double.IsNaN(scale) || double.IsInfinity(scale))
+    {
+        scale = 1.0;
+    }
     return gradient.Multiply(NumOps.FromDouble(scale));
 }

src/MixedPrecision/Float8Types.cs (3)

82-86: Unused constants can be removed.

ExponentBits (line 82) and ExponentMask (line 85) are declared but never used in this struct. Consider removing them to reduce code clutter, or use them in place of hardcoded values (e.g., 0xF on line 228).

---

311-315: Same unused constants as in Float8E4M3.

ExponentBits and ExponentMask are also unused here. Apply the same cleanup if addressing this in Float8E4M3.

---

550-558: Add null validation for extension methods.

These extension methods will throw NullReferenceException if passed a null array. Consider adding explicit null checks to throw ArgumentNullException with a meaningful parameter name for better API usability.

♻️ Example fix for ToE4M3

 public static Float8E4M3[] ToE4M3(this float[] values)
 {
+    ArgumentNullException.ThrowIfNull(values);
     var result = new Float8E4M3[values.Length];
     for (int i = 0; i < values.Length; i++)
     {
         result[i] = Float8E4M3.FromFloat(values[i]);
     }
     return result;
 }

Apply the same pattern to ToE5M2, and both ToFloatArray overloads.

[coderabbitai]

Actionable comments posted: 1

🤖 Fix all issues with AI agents

In `@src/MixedPrecision/Float8Types.cs`:
- Around line 291-327: The MinPositive field and XML docs for Float8E5M2 are
incorrect: FromFloat treats 2^-16 as the smallest subnormal but the remarks and
MinPositive currently reflect the smallest normal (2^-14). Update the XML
remarks to state the smallest positive subnormal is 2^-16 (~1.5258789e-5) and
change the MinPositive initializer to FromFloat(1.52587890625e-5f) (or compute
as MathF.Pow(2f, -16)) so the documentation and the Float8E5M2.MinPositive value
match FromFloat's behavior; ensure the remarks language clarifies "smallest
positive (subnormal)" where appropriate.

[Copilot]

Pull request overview

Copilot reviewed 3 out of 3 changed files in this pull request and generated no new comments.

---

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

[sonarqubecloud]

Quality Gate failed

Failed conditions
0.0% Coverage on New Code (required ≥ 80%)

See analysis details on SonarQube Cloud