diff --git a/src/data/loader.rs b/src/data/loader.rs index 366cdbf..dc19063 100644 --- a/src/data/loader.rs +++ b/src/data/loader.rs @@ -1,7 +1,7 @@ //! Batching and shuffling. use crate::rng; -use crate::tensor::{Device, Tensor}; +use crate::tensor::{ops::index::index_from_f32, Device, Tensor}; use super::dataset::Dataset; @@ -22,8 +22,14 @@ impl Batch { } /// Targets read as class indices — the form [`cross_entropy`](crate::nn::cross_entropy) takes. + /// + /// Panics if a target is not a finite, non-negative integer representable as `usize`. pub fn labels(&self) -> Vec { - self.targets.to_vec().iter().map(|&v| v as usize).collect() + self.targets + .to_vec() + .into_iter() + .map(index_from_f32) + .collect() } } diff --git a/src/nn/embedding.rs b/src/nn/embedding.rs index a399579..8a0ff20 100644 --- a/src/nn/embedding.rs +++ b/src/nn/embedding.rs @@ -1,6 +1,6 @@ //! Embedding table and positional encoding. -use crate::tensor::Tensor; +use crate::tensor::{ops::index::index_from_f32, Tensor}; use super::module::Module; use super::param::Param; @@ -54,10 +54,11 @@ impl Embedding { } impl Module for Embedding { - /// Treats the input's values as ids. Prefer [`lookup`](Embedding::lookup) — - /// it takes `usize` and cannot silently truncate a float. + /// Treats finite, non-negative integer values as ids. + /// Panics if an id cannot be represented as `usize`. + /// Prefer [`lookup`](Embedding::lookup), which takes `usize` directly. fn forward(&self, input: &Tensor) -> Tensor { - let ids: Vec = input.to_vec().iter().map(|&v| v as usize).collect(); + let ids: Vec = input.to_vec().into_iter().map(index_from_f32).collect(); let mut shape: Vec = input.shape().iter().map(|&d| d as i64).collect(); shape.push(self.dim as i64); self.lookup(&ids).reshape(&shape) diff --git a/src/tensor/ops/index.rs b/src/tensor/ops/index.rs index e945e7e..b6bad10 100644 --- a/src/tensor/ops/index.rs +++ b/src/tensor/ops/index.rs @@ -8,6 +8,18 @@ use crate::cuda::kernels; use crate::tensor::storage::Storage; use crate::tensor::Tensor; +/// Convert a float index without truncating or saturating it. +pub(crate) fn index_from_f32(value: f32) -> usize { + assert!( + value.is_finite() + && value >= 0.0 + && value.fract() == 0.0 + && f64::from(value) < usize::MAX as f64 + 1.0, + "invalid index {value}: expected a finite, non-negative integer that fits usize" + ); + value as usize +} + impl Tensor { /// Gather rows `ids` from this `[rows, cols]` table into `[ids.len(), cols]`. /// diff --git a/tests/robustness.rs b/tests/robustness.rs index 842e1de..c8fb778 100644 --- a/tests/robustness.rs +++ b/tests/robustness.rs @@ -56,6 +56,96 @@ fn train_steps( } } +// ── Float indices ──────────────────────────────────────────────────────────── + +fn assert_invalid_float_indices(mut consume: impl FnMut(f32)) { + let upper = (usize::MAX as f64 + 1.0) as f32; + let mut failures = Vec::new(); + for value in [ + -1.0, + f32::NAN, + 2.5, + f32::INFINITY, + f32::NEG_INFINITY, + upper, + f32::MAX, + ] { + let caught = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| consume(value))); + match caught { + Ok(()) => failures.push(format!("accepted {value}")), + Err(payload) => { + let message = payload + .downcast_ref::() + .map(String::as_str) + .or_else(|| payload.downcast_ref::<&str>().copied()) + .unwrap_or("non-string panic"); + if !message.contains(&value.to_string()) { + failures.push(format!("{value} missing from panic: {message}")); + } + } + } + } + assert!(failures.is_empty(), "{}", failures.join("; ")); +} + +#[test] +fn float_indices_embedding_rejects_invalid_values() { + let embedding = Embedding::new(4, 2); + assert_invalid_float_indices(|value| { + embedding.forward(&Tensor::from_vec(vec![value], &[1])); + }); +} + +#[test] +fn float_indices_labels_reject_invalid_values() { + assert_invalid_float_indices(|value| { + Batch { + inputs: Tensor::zeros(&[1, 1]), + targets: Tensor::from_vec(vec![value], &[1]), + } + .labels(); + }); +} + +#[test] +fn float_indices_embedding_preserves_valid_ids() { + let embedding = Embedding::new(4, 2); + let ids = Tensor::from_vec(vec![0.0, 3.0, -0.0, 2.0], &[2, 2]); + let output = embedding.forward(&ids); + assert_eq!(output.shape(), &[2, 2, 2]); + assert_eq!(output.to_vec(), embedding.lookup(&[0, 3, 0, 2]).to_vec()); +} + +#[test] +fn float_indices_labels_preserve_valid_ids() { + let batch = Batch { + inputs: Tensor::zeros(&[4, 1]), + targets: Tensor::from_vec(vec![0.0, 3.0, -0.0, 2.0], &[4]), + }; + assert_eq!(batch.labels(), vec![0, 3, 0, 2]); +} + +#[test] +#[cfg(any(target_pointer_width = "32", target_pointer_width = "64"))] +fn float_indices_labels_accept_large_integers() { + let batch = Batch { + inputs: Tensor::zeros(&[1, 1]), + targets: Tensor::from_vec(vec![16_777_218.0], &[1]), + }; + assert_eq!(batch.labels(), vec![16_777_218]); +} + +#[test] +fn float_indices_labels_accept_largest_fitting_integer() { + let upper = (usize::MAX as f64 + 1.0) as f32; + let largest = f32::from_bits(upper.to_bits() - 1).floor(); + let batch = Batch { + inputs: Tensor::zeros(&[1, 1]), + targets: Tensor::from_vec(vec![largest], &[1]), + }; + assert_eq!(batch.labels(), vec![largest as usize]); +} + // ── Resuming a run ─────────────────────────────────────────────────────────── #[test]