diff --git a/src/nn/linear.rs b/src/nn/linear.rs index 5673e26..1bd45dc 100644 --- a/src/nn/linear.rs +++ b/src/nn/linear.rs @@ -94,6 +94,7 @@ impl Module for Linear { /// Uniform over `±1/√fan_in`, the bias init shared with the conv layers. pub(super) fn uniform_bias(out_features: usize, fan_in: usize) -> Tensor { + assert!(fan_in > 0, "uniform_bias: fan_in must be positive, got 0"); let bound = 1.0 / (fan_in as f32).sqrt(); Tensor::uniform(&[out_features], -bound, bound) } diff --git a/src/rng.rs b/src/rng.rs index 7d28b86..de65a4d 100644 --- a/src/rng.rs +++ b/src/rng.rs @@ -22,6 +22,9 @@ pub fn with_rng(f: impl FnOnce(&mut StdRng) -> T) -> T { /// `n` samples from the uniform distribution on `[lo, hi)`. pub fn uniform(n: usize, lo: f32, hi: f32) -> Vec { + if lo == hi { + return vec![lo; n]; + } with_rng(|rng| (0..n).map(|_| rng.random_range(lo..hi)).collect()) } diff --git a/src/tensor/init.rs b/src/tensor/init.rs index d2115fe..982e497 100644 --- a/src/tensor/init.rs +++ b/src/tensor/init.rs @@ -43,12 +43,20 @@ impl Tensor { /// Kaiming (He) uniform, bounds `±√(6 / fan_in)`. For ReLU-family layers. pub fn kaiming_uniform(shape: &[usize], fan_in: usize) -> Tensor { + assert!( + fan_in > 0, + "kaiming_uniform: fan_in must be positive, got 0 for shape {shape:?}" + ); let bound = (6.0 / fan_in as f32).sqrt(); Tensor::uniform(shape, -bound, bound) } /// Xavier (Glorot) uniform, for `tanh` and `sigmoid` layers. pub fn xavier_uniform(shape: &[usize], fan_in: usize, fan_out: usize) -> Tensor { + assert!( + fan_in + fan_out > 0, + "xavier_uniform: fan_in + fan_out must be positive, got 0 for shape {shape:?}" + ); let bound = (6.0 / (fan_in + fan_out) as f32).sqrt(); Tensor::uniform(shape, -bound, bound) } diff --git a/tests/linear.rs b/tests/linear.rs index 93edbe6..02a8020 100644 --- a/tests/linear.rs +++ b/tests/linear.rs @@ -63,3 +63,27 @@ fn no_bias_preserves_forward_shape_and_trainable_weight() { output.sum().backward(); assert_eq!(layer.weight.grad().unwrap().shape(), &[3, 2]); } + +#[test] +#[should_panic(expected = "fan_in must be positive")] +fn linear_zero_in_features_panics_with_clear_message() { + let _ = Linear::new(0, 10); +} + +#[test] +#[should_panic(expected = "fan_in must be positive")] +fn linear_no_bias_zero_in_features_panics() { + let _ = Linear::no_bias(0, 10); +} + +#[test] +#[should_panic(expected = "fan_in + fan_out must be positive")] +fn xavier_zero_fan_panics() { + let _ = Tensor::xavier_uniform(&[3, 3], 0, 0); +} + +#[test] +fn rng_uniform_degenerate_range_returns_constant() { + let vals = fastnn::rng::uniform(5, 3.0, 3.0); + assert_eq!(vals, vec![3.0; 5]); +}