feat(rtx-nn): f64-capable GenericLinear + f64 MLP gradient-precision capstone

Phase 4/5 of the rustytorch f32→f64 plan. GenericLinear's 3
`impl<B: Backend<FloatElem = f32>>` blocks relaxed to `impl<B: Backend>`
(from_weights takes &[B::FloatElem]; Xavier scale via B::FloatElem::from_f64), so a
Linear→ReLU→Linear MLP runs end-to-end on CpuBackendF64. f32 backward-compat holds
via B::FloatElem = f32.

Capstone (tests/f64_mlp_precision.rs): a 4→8→1 MLP gradient checked vs central
finite differences — f64 err 6.99e-12 (≤1e-9) vs f32 err 1.01e-2, i.e. f64 ~1.45e9×
more accurate. This is the quantum-precision-gradient win that motivated the migration.

Validated: rtx-nn 334 f32 lib tests + 2 new f64 capstone tests pass; rtx-autograd
builds + tests pass; **QPUDIDP qpu-didp-surrogate compiles + 15 tests pass** (uses
GenericLinear). clippy clean.

Scope note: rtx-autograd's reverse-mode tape stores f32 concretely
(backward()->HashMap<_,Vec<f32>>) — making it f64 is a deep tape re-architecture, not
a constraint relaxation, so it's a documented follow-on (no current consumer uses it;
QPUDIDP hand-rolls f64 backprop). Other rtx-nn layers (conv/transformer/attention/...)
remain f32-gated — same mechanical relaxation, follow-on.

Co-Authored-By: Claude Opus 4.8 (1M context) <[email protected]>
This commit is contained in:
Claude Code
2026-06-26 22:36:23 -07:00
co-authored by Claude Opus 4.8
parent 877071bb9a
commit c28a848250
2 changed files with 140 additions and 7 deletions
+7 -7
View File
@@ -2,7 +2,7 @@
//! //!
//! Provides a fully-connected (dense) layer that works with any backend. //! Provides a fully-connected (dense) layer that works with any backend.
use rtx_backend::Backend; use rtx_backend::{Backend, FloatElement};
use rtx_tensor::generic::GenericTensor; use rtx_tensor::generic::GenericTensor;
use std::fmt::{self, Debug}; use std::fmt::{self, Debug};
@@ -44,7 +44,7 @@ pub struct GenericLinear<B: Backend> {
training: bool, training: bool,
} }
impl<B: Backend<FloatElem = f32>> GenericLinear<B> { impl<B: Backend> GenericLinear<B> {
/// Create a new Linear layer. /// Create a new Linear layer.
/// ///
/// # Arguments /// # Arguments
@@ -57,7 +57,7 @@ impl<B: Backend<FloatElem = f32>> GenericLinear<B> {
// Xavier uniform initialization: U(-sqrt(6/(in+out)), sqrt(6/(in+out))) // Xavier uniform initialization: U(-sqrt(6/(in+out)), sqrt(6/(in+out)))
// For simplicity, using randn scaled by sqrt(2/in_features) // For simplicity, using randn scaled by sqrt(2/in_features)
let weight = GenericTensor::randn([out_features, in_features], device); let weight = GenericTensor::randn([out_features, in_features], device);
let scale = (2.0 / in_features as f32).sqrt(); let scale = B::FloatElem::from_f64((2.0 / in_features as f64).sqrt());
let weight = weight.mul_scalar(scale); let weight = weight.mul_scalar(scale);
let bias = if bias { let bias = if bias {
@@ -87,8 +87,8 @@ impl<B: Backend<FloatElem = f32>> GenericLinear<B> {
/// * `out_features` - Size of each output sample /// * `out_features` - Size of each output sample
/// * `device` - Device to create tensors on /// * `device` - Device to create tensors on
pub fn from_weights( pub fn from_weights(
weight_data: &[f32], weight_data: &[B::FloatElem],
bias_data: Option<&[f32]>, bias_data: Option<&[B::FloatElem]>,
in_features: usize, in_features: usize,
out_features: usize, out_features: usize,
device: &B::Device, device: &B::Device,
@@ -144,7 +144,7 @@ impl<B: Backend<FloatElem = f32>> GenericLinear<B> {
} }
} }
impl<B: Backend<FloatElem = f32>> GenericModule<B> for GenericLinear<B> { impl<B: Backend> GenericModule<B> for GenericLinear<B> {
fn forward(&self, input: &GenericTensor<B, 2>) -> GenericTensor<B, 2> { fn forward(&self, input: &GenericTensor<B, 2>) -> GenericTensor<B, 2> {
// y = x @ W^T + b // y = x @ W^T + b
// input: [batch, in_features] // input: [batch, in_features]
@@ -196,7 +196,7 @@ impl<B: Backend<FloatElem = f32>> GenericModule<B> for GenericLinear<B> {
} }
// Special implementation to return actual parameters // Special implementation to return actual parameters
impl<B: Backend<FloatElem = f32>> GenericLinear<B> { impl<B: Backend> GenericLinear<B> {
/// Get all parameters as 2D tensors. /// Get all parameters as 2D tensors.
/// ///
/// Returns weight and (reshaped) bias tensors. /// Returns weight and (reshaped) bias tensors.
@@ -0,0 +1,133 @@
//! Capstone for the rustytorch f32→f64 migration (Phase 4/5).
//!
//! Builds the same MLP — `Linear(4→8) → ReLU → Linear(8→1)` — on the new
//! `CpuBackendF64` and the original f32 `CpuBackend`, and shows that the f64 path
//! computes a gradient (via central finite differences over the genericized
//! `GenericLinear`/`GenericTensor` forward) that matches the analytic gradient to
//! ≤1e-9, whereas the f32 path is roundoff-limited to ~1e-2 — the precision win
//! that motivates the migration.
//!
//! (Reverse-mode `rtx-autograd` is still f32-concrete — its tape returns
//! `Vec<f32>` — so this validates the f64 *layer* path with an analytic backward,
//! per the plan's documented fallback.)
use rtx_backend::{Backend, FloatElement};
use rtx_backend_cpu::{CpuBackend, CpuBackendF64, CpuDevice};
use rtx_nn::{GenericLinear, GenericModule};
use rtx_tensor::generic::GenericTensor;
const IN: usize = 4;
const HID: usize = 8;
const T: f64 = 0.5; // regression target
// Nominal f64 parameters, chosen so h_pre[0] > 0 (ReLU active, smooth gradient).
fn w1() -> Vec<f64> {
(0..HID * IN).map(|k| 0.1 + 0.01 * k as f64).collect()
}
fn b1() -> Vec<f64> {
(0..HID).map(|i| 0.2 + 0.01 * i as f64).collect()
}
fn w2() -> Vec<f64> {
(0..HID).map(|j| 0.15 + 0.01 * j as f64).collect()
}
fn b2() -> Vec<f64> {
vec![0.1]
}
fn xin() -> Vec<f64> {
vec![0.5, 0.3, 0.8, 0.2]
}
fn cast<B: Backend>(v: &[f64]) -> Vec<B::FloatElem> {
v.iter().map(|&x| B::FloatElem::from_f64(x)).collect()
}
/// Loss `L = 0.5·(y T)²` of the MLP forward, computed at the precision of `B`
/// (weights passed as `B::FloatElem`). Exercises the genericized GenericLinear +
/// GenericTensor::relu under the chosen backend.
fn mlp_loss<B: Backend>(w1: &[f64], dev: &B::Device) -> f64 {
let l1 = GenericLinear::<B>::from_weights(
&cast::<B>(w1),
Some(&cast::<B>(&b1())),
IN,
HID,
dev,
);
let l2 = GenericLinear::<B>::from_weights(
&cast::<B>(&w2()),
Some(&cast::<B>(&b2())),
HID,
1,
dev,
);
let x = GenericTensor::<B, 2>::from_slice(&cast::<B>(&xin()), [1, IN], dev);
let h = l1.forward(&x).relu();
let y = l2.forward(&h);
let yv = y.to_vec()[0].to_f64();
0.5 * (yv - T) * (yv - T)
}
/// Central-difference dL/dW1[0,0] at precision `B`, step `eps`.
fn fd_grad_w1_00<B: Backend>(eps: f64, dev: &B::Device) -> f64 {
let mut wp = w1();
wp[0] += eps;
let mut wm = w1();
wm[0] -= eps;
(mlp_loss::<B>(&wp, dev) - mlp_loss::<B>(&wm, dev)) / (2.0 * eps)
}
/// Analytic dL/dW1[0,0] via manual f64 backprop through the same MLP.
fn analytic_grad_w1_00() -> f64 {
let (w1, b1, w2, b2, x) = (w1(), b1(), w2(), b2(), xin());
let mut h_pre = vec![0.0f64; HID];
for i in 0..HID {
let mut s = b1[i];
for j in 0..IN {
s += w1[i * IN + j] * x[j];
}
h_pre[i] = s;
}
let h: Vec<f64> = h_pre.iter().map(|&v| v.max(0.0)).collect();
let mut y = b2[0];
for i in 0..HID {
y += w2[i] * h[i];
}
let dy = y - T;
let relu_p = if h_pre[0] > 0.0 { 1.0 } else { 0.0 };
// dL/dW1[0,0] = (yT)·W2[0]·relu'(h_pre[0])·x[0]
dy * w2[0] * relu_p * x[0]
}
#[test]
fn f64_mlp_gradient_matches_finite_difference_to_1e9() {
let dev = CpuDevice::new();
let g_an = analytic_grad_w1_00();
assert!(g_an.abs() > 1e-6, "sanity: gradient should be non-trivial ({g_an:.3e})");
let g_fd = fd_grad_w1_00::<CpuBackendF64>(1e-6, &dev);
let err = (g_fd - g_an).abs();
println!("f64: analytic={g_an:.15} fd={g_fd:.15} err={err:.3e}");
assert!(
err <= 1e-9,
"f64 finite-diff gradient must match analytic to 1e-9, got {err:.3e}"
);
}
#[test]
fn f64_beats_f32_gradient_precision() {
let dev = CpuDevice::new();
let g_an = analytic_grad_w1_00();
let err_f64 = (fd_grad_w1_00::<CpuBackendF64>(1e-6, &dev) - g_an).abs();
let err_f32 = (fd_grad_w1_00::<CpuBackend>(1e-6, &dev) - g_an).abs();
println!(
"err_f64={err_f64:.3e} err_f32={err_f32:.3e} ratio={:.0}x",
err_f32 / err_f64.max(1e-300)
);
assert!(err_f64 <= 1e-9, "f64 err {err_f64:.3e}");
assert!(
err_f32 > 1e-4,
"f32 finite-diff should be roundoff-limited (>1e-4), got {err_f32:.3e}"
);
assert!(
err_f64 < err_f32 * 1e-3,
"f64 must be >=1000x more accurate (f64={err_f64:.3e}, f32={err_f32:.3e})"
);
}