reject invalid float indices before conversion

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<usize> {
-        self.targets.to_vec().iter().map(|&v| v as usize).collect()
+        self.targets
+            .to_vec()
+            .into_iter()
+            .map(index_from_f32)
+            .collect()
     }
 }
 

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<usize> = input.to_vec().iter().map(|&v| v as usize).collect();
+        let ids: Vec<usize> = input.to_vec().into_iter().map(index_from_f32).collect();
         let mut shape: Vec<i64> = input.shape().iter().map(|&d| d as i64).collect();
         shape.push(self.dim as i64);
         self.lookup(&ids).reshape(&shape)

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]`.
     ///

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::<String>()
+                    .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]