diff --git a/src/nn/rnn.rs b/src/nn/rnn.rs index 3426a27..f36cb10 100644 --- a/src/nn/rnn.rs +++ b/src/nn/rnn.rs @@ -185,7 +185,13 @@ fn check_sequence(input: &Tensor, input_size: usize, layer: &str) -> (usize, usi "{layer} expects {input_size} features, got {:?}", input.shape() ); - (input.dim(0), input.dim(1)) + let (batch, steps) = (input.dim(0), input.dim(1)); + assert!( + steps > 0, + "{layer}: sequence length is 0 in shape {:?}", + input.shape() + ); + (batch, steps) } /// Timestep `t` as `[batch, features]`. diff --git a/src/tensor/ops/view.rs b/src/tensor/ops/view.rs index cefde89..f64d69d 100644 --- a/src/tensor/ops/view.rs +++ b/src/tensor/ops/view.rs @@ -200,6 +200,7 @@ impl Tensor { /// Join tensors along a new dimension inserted at `axis`. pub fn stack(parts: &[&Tensor], axis: usize) -> Tensor { + assert!(!parts.is_empty(), "stack: needs at least one tensor"); let lifted: Vec = parts.iter().map(|t| t.unsqueeze(axis)).collect(); Tensor::cat(&lifted.iter().collect::>(), axis) } diff --git a/tests/rnn.rs b/tests/rnn.rs new file mode 100644 index 0000000..7a8c284 --- /dev/null +++ b/tests/rnn.rs @@ -0,0 +1,21 @@ +use fastnn::prelude::*; + +#[test] +#[should_panic(expected = "sequence length is 0")] +fn lstm_zero_length_sequence_panics() { + let lstm = LSTM::new(2, 3); + lstm.forward(&Tensor::zeros(&[1, 0, 2])); +} + +#[test] +#[should_panic(expected = "sequence length is 0")] +fn gru_zero_length_sequence_panics() { + let gru = GRU::new(2, 3); + gru.forward(&Tensor::zeros(&[1, 0, 2])); +} + +#[test] +#[should_panic(expected = "stack: needs at least one tensor")] +fn stack_empty_panics() { + Tensor::stack(&[], 0); +}