Files
foxhunt/crates/ml-core/src/xavier_init.rs
jgrusewski 22004a7368 refactor(cuda): eliminate candle from ml-core, ml-ppo, and 4 thin crates
Hard refactor — no shims, no compat layers. Candle removed from Cargo.toml
and all source files in 6 crates:

- ml-core: MlDevice enum, checkpoint.rs (safetensors direct), cudarc imports
  fixed from candle re-export to direct, AdamWConfig lr_decay, cuda_compat
  gutted. Net -7,341 lines.
- ml-ppo: All 16 files rewritten. LSTM→CudaLSTM, VarMap→GpuVarStore,
  PPOAgent 2306→700 lines, checkpoint→binary format.
- ml-ensemble: GPU-resident sigmoid via custom CUDA kernel.
- ml-explainability: Integrated gradients via GPU finite-difference kernels.
- ml-labeling: Device→MlDevice.
- ml-hyperopt: Cargo.toml only.

Remaining: ml-dqn (24 files), ml-supervised (4 files), ml crate (104 files).

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
2026-03-17 22:27:56 +01:00

103 lines
3.2 KiB
Rust

//! Xavier/Glorot initialization helpers (CPU-side).
//!
//! GPU-side initialization is handled by `cuda_autograd::init` which directly
//! allocates and fills GPU buffers. This module provides CPU-side convenience
//! functions used by downstream crates that still operate on host data.
//!
//! Reference: Glorot & Bengio (2010)
/// Compute the Xavier uniform limit: `sqrt(6 / (fan_in + fan_out))`.
pub fn xavier_limit(fan_in: usize, fan_out: usize) -> f64 {
(6.0 / (fan_in + fan_out) as f64).sqrt()
}
/// Compute the Kaiming uniform limit: `sqrt(6 / fan_in)`.
pub fn kaiming_limit(fan_in: usize) -> f64 {
(6.0 / fan_in as f64).sqrt()
}
/// Generate `n` uniform random f32 values in `[lo, hi)`.
///
/// Uses a simple xoshiro256++ PRNG seeded from system time + thread id.
pub fn generate_uniform(n: usize, lo: f64, hi: f64) -> Vec<f32> {
use std::collections::hash_map::DefaultHasher;
use std::hash::{Hash, Hasher};
use std::time::SystemTime;
let mut hasher = DefaultHasher::new();
SystemTime::now().hash(&mut hasher);
std::thread::current().id().hash(&mut hasher);
let seed = hasher.finish();
let mut s = [
seed,
seed.wrapping_mul(6364136223846793005).wrapping_add(1),
seed.wrapping_mul(1442695040888963407).wrapping_add(3),
seed.wrapping_mul(2891336453).wrapping_add(7),
];
let range = hi - lo;
let mut out = Vec::with_capacity(n);
for _ in 0..n {
let result = s[0].wrapping_add(s[3]).rotate_left(23).wrapping_add(s[0]);
let t = s[1] << 17;
s[2] ^= s[0];
s[3] ^= s[1];
s[1] ^= s[2];
s[0] ^= s[3];
s[2] ^= t;
s[3] = s[3].rotate_left(45);
let f = (result >> 11) as f64 / (1_u64 << 53) as f64;
out.push((lo + f * range) as f32);
}
out
}
/// Generate Xavier uniform weights on CPU: shape `[fan_out, fan_in]`.
pub fn xavier_uniform_cpu(fan_in: usize, fan_out: usize) -> Vec<f32> {
let limit = xavier_limit(fan_in, fan_out);
generate_uniform(fan_out * fan_in, -limit, limit)
}
/// Verify Xavier initialization statistics for testing.
///
/// Returns `(mean, variance, expected_variance)`.
pub fn verify_xavier_stats_cpu(
weights: &[f32],
fan_in: usize,
fan_out: usize,
) -> (f32, f32, f32) {
let n = weights.len() as f32;
let mean: f32 = weights.iter().sum::<f32>() / n;
let variance: f32 = weights.iter().map(|x| (x - mean) * (x - mean)).sum::<f32>() / n;
let expected_var = 2.0 / (fan_in + fan_out) as f32;
(mean, variance, expected_var)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_xavier_uniform_cpu_shape() {
let weights = xavier_uniform_cpu(64, 128);
assert_eq!(weights.len(), 128 * 64);
}
#[test]
fn test_xavier_uniform_cpu_statistics() {
let fan_in = 256;
let fan_out = 256;
let weights = xavier_uniform_cpu(fan_in, fan_out);
let (mean, variance, expected_var) = verify_xavier_stats_cpu(&weights, fan_in, fan_out);
assert!(mean.abs() < 0.02, "Mean should be ~0, got {mean}");
assert!(
(variance - expected_var).abs() / expected_var < 0.2,
"Variance {variance} should be ~{expected_var}"
);
}
}