Function test_decoder_block_forward

cake-core/src/models/vibevoice/vae_decoder.rs:528–560

()

    #[test]
    fn test_decoder_block_forward() {
        let ch = 32;
        let mut map = std::collections::HashMap::new();
        let dev = Device::Cpu;

        fn mt(shape: &[usize], seed: u64) -> Tensor {
            use rand::{Rng, SeedableRng};
            let n: usize = shape.iter().product();
            let mut rng = rand::rngs::StdRng::seed_from_u64(seed);
            let d: Vec<f32> = (0..n).map(|_| rng.gen_range(-0.01..0.01)).collect();
            Tensor::from_vec(d, shape, &Device::Cpu).unwrap()
        }

        map.insert("norm.weight".into(), Tensor::ones(ch, DType::F32, &dev).unwrap());
        map.insert("gamma".into(), Tensor::ones(ch, DType::F32, &dev).unwrap());
        map.insert("mixer.conv.conv.conv.weight".into(), mt(&[ch, 1, 7], 1));
        map.insert("mixer.conv.conv.conv.bias".into(), mt(&[ch], 2));
        map.insert("ffn_norm.weight".into(), Tensor::ones(ch, DType::F32, &dev).unwrap());
        map.insert("ffn_gamma".into(), Tensor::ones(ch, DType::F32, &dev).unwrap());
        map.insert("ffn.linear1.weight".into(), mt(&[ch * 4, ch], 3));
        map.insert("ffn.linear1.bias".into(), mt(&[ch * 4], 4));
        map.insert("ffn.linear2.weight".into(), mt(&[ch, ch * 4], 5));
        map.insert("ffn.linear2.bias".into(), mt(&[ch], 6));

        let vb = VarBuilder::from_tensors(map, DType::F32, &dev);
        let backend = crate::backends::create_backend(&dev);
        let block = DecoderBlock::load(vb, ch, 1e-5, backend).unwrap();

        // Input: (batch=1, channels=32, seq=16)
        let x = mt(&[1, ch, 16], 10);
        let y = block.forward(&x).unwrap();
        assert_eq!(y.dims(), &[1, ch, 16]);
    }

    // --- StreamingConvCache ---