Test the reset too.

huggingface · Sep 22, 2024 · f9579f8 · f9579f8
1 parent 1bddd44
commit f9579f8
Showing 1 changed file with 19 additions and 14 deletions.
diff --git a/candle-nn/tests/kv_cache.rs b/candle-nn/tests/kv_cache.rs
@@ -9,19 +9,24 @@ use candle::{Device, Result, Tensor};
 #[test]
 fn kv_cache() -> Result<()> {
  let mut cache = candle_nn::kv_cache::Cache::new(0, 16);
- let data = cache.current_data()?;
- assert!(data.is_none());
- let t = Tensor::new(&[1f32, 2., 3.], &Device::Cpu)?;
- cache.append(&t)?;
- let data = cache.current_data()?.unwrap();
- assert_eq!(data.to_vec1::<f32>()?, [1., 2., 3.]);
- let t = Tensor::new(&[4f32], &Device::Cpu)?;
- cache.append(&t)?;
- let data = cache.current_data()?.unwrap();
- assert_eq!(data.to_vec1::<f32>()?, [1., 2., 3., 4.]);
- let t = Tensor::new(&[0f32, 5., 6., 7.], &Device::Cpu)?;
- cache.append(&t)?;
- let data = cache.current_data()?.unwrap();
- assert_eq!(data.to_vec1::<f32>()?, [1., 2., 3., 4., 0., 5., 6., 7.]);
+ for _ in [0, 1] {
+ assert_eq!(cache.current_seq_len(), 0);
+ let data = cache.current_data()?;
+ assert!(data.is_none());
+ let t = Tensor::new(&[1f32, 2., 3.], &Device::Cpu)?;
+ cache.append(&t)?;
+ let data = cache.current_data()?.unwrap();
+ assert_eq!(data.to_vec1::<f32>()?, [1., 2., 3.]);
+ let t = Tensor::new(&[4f32], &Device::Cpu)?;
+ cache.append(&t)?;
+ let data = cache.current_data()?.unwrap();
+ assert_eq!(data.to_vec1::<f32>()?, [1., 2., 3., 4.]);
+ let t = Tensor::new(&[0f32, 5., 6., 7.], &Device::Cpu)?;
+ cache.append(&t)?;
+ let data = cache.current_data()?.unwrap();
+ assert_eq!(data.to_vec1::<f32>()?, [1., 2., 3., 4., 0., 5., 6., 7.]);
+ assert_eq!(cache.current_seq_len(), 8);
+ cache.reset();
+ }
  Ok(())
 }