extern crate alloc;
use burn::prelude::*;
use burn::nn::Linear;
use burn::nn::LinearConfig;
use burn::nn::Lstm;
use burn::nn::LstmConfig;
use burn::nn::LstmState;
#[derive(Module, Debug)]
pub struct Model {
constant1: burn::module::Param<Tensor<2>>,
lstm1: Lstm,
lstm2: Lstm,
linear1: Linear,
linear2: Linear,
linear3: Linear,
#[module(skip)]
device: Device,
}
impl Model {
#[allow(unused_variables)]
pub fn new(device: &Device) -> Self {
let constant1: burn::module::Param<Tensor<2>> = burn::module::Param::uninitialized(
burn::module::ParamId::new(),
move |device, _require_grad| Tensor::<
2,
>::zeros([13088, 640], (device, burn::tensor::DType::F32)),
device.clone(),
false,
[13088, 640].into(),
);
let lstm1 = LstmConfig::new(640, 640, true)
.with_batch_first(false)
.with_input_forget(false)
.init(device);
let lstm2 = LstmConfig::new(640, 640, true)
.with_batch_first(false)
.with_input_forget(false)
.init(device);
let linear1 = LinearConfig::new(1024, 640).with_bias(true).init(device);
let linear2 = LinearConfig::new(640, 640).with_bias(true).init(device);
let linear3 = LinearConfig::new(640, 13088).with_bias(true).init(device);
Self {
constant1,
lstm1,
lstm2,
linear1,
linear2,
linear3,
device: device.clone(),
}
}
#[allow(clippy::let_and_return, clippy::approx_constant)]
pub fn forward(
&self,
encoder_outputs: Tensor<3>,
targets: Tensor<2, Int>,
target_length: Tensor<1, Int>,
input_states_1: Tensor<3>,
input_states_2: Tensor<3>,
) -> (Tensor<4>, Tensor<1, Int>, Tensor<3>, Tensor<3>) {
let constant1_out1 = self.constant1.val();
let cast1_out1 = targets.cast(burn::tensor::DType::I64);
let gather1_out1 = {
let axis_size = constant1_out1.dims()[0] as i64;
let negative = cast1_out1.clone().lower_elem(0i64);
let corrected = cast1_out1.clone() + axis_size;
let indices = cast1_out1.mask_where(negative, corrected);
constant1_out1.take::<2, 3>(0, indices)
};
let transpose1_out1 = gather1_out1.permute([1, 0, 2]);
let slice1_out1 = input_states_1.clone().slice(s![0..1, .., ..]);
let slice2_out1 = input_states_2.clone().slice(s![0..1, .., ..]);
let (lstm1_out1, lstm1_out2, lstm1_out3) = {
let (output_seq, final_state) = self
.lstm1
.forward(
transpose1_out1,
Some(
LstmState::new(
slice2_out1.squeeze_dim(0),
slice1_out1.squeeze_dim(0),
),
),
);
(
output_seq.unsqueeze_dims::<4>(&[1]),
final_state.hidden.unsqueeze_dims::<3>(&[0]),
final_state.cell.unsqueeze_dims::<3>(&[0]),
)
};
let transpose2_out1 = lstm1_out1.permute([0, 2, 1, 3]);
let shape1_out1: [i64; 4] = {
let axes = &transpose2_out1.clone().dims()[0..4];
let mut output = [0i64; 4];
for i in 0..4 {
output[i] = axes[i] as i64;
}
output
};
let slice3_out1: [i64; 1] = shape1_out1[0..1].try_into().unwrap();
let slice4_out1: [i64; 1] = shape1_out1[1..2].try_into().unwrap();
let slice5_out1: [i64; 1] = shape1_out1[2..3].try_into().unwrap();
let slice6_out1: [i64; 1] = shape1_out1[3..4].try_into().unwrap();
let mul1_out1 = {
let __lhs = slice5_out1;
let __rhs = slice6_out1;
core::array::from_fn::<
i64,
1usize,
_,
>(|__i| (__lhs[__i]).saturating_mul(__rhs[__i]))
};
let reshape1_out1 = mul1_out1;
let concat1_out1: [i64; 3usize] = [
&slice3_out1[..],
&slice4_out1[..],
&reshape1_out1[..],
]
.concat()
.try_into()
.unwrap();
let reshape2_out1 = transpose2_out1.reshape(concat1_out1);
let slice7_out1 = input_states_1.slice(s![1..2, .., ..]);
let slice8_out1 = input_states_2.slice(s![1..2, .., ..]);
let (lstm2_out1, lstm2_out2, lstm2_out3) = {
let (output_seq, final_state) = self
.lstm2
.forward(
reshape2_out1,
Some(
LstmState::new(
slice8_out1.squeeze_dim(0),
slice7_out1.squeeze_dim(0),
),
),
);
(
output_seq.unsqueeze_dims::<4>(&[1]),
final_state.hidden.unsqueeze_dims::<3>(&[0]),
final_state.cell.unsqueeze_dims::<3>(&[0]),
)
};
let transpose3_out1 = lstm2_out1.permute([0, 2, 1, 3]);
let reshape3_out1 = transpose3_out1.reshape([1, -1, 640]);
let concat2_out1 = burn::tensor::Tensor::cat([lstm1_out2, lstm2_out2].into(), 0);
let concat3_out1 = burn::tensor::Tensor::cat([lstm1_out3, lstm2_out3].into(), 0);
let transpose4_out1 = encoder_outputs.permute([0, 2, 1]);
let transpose5_out1 = reshape3_out1.permute([1, 0, 2]);
let linear1_out1 = self.linear1.forward(transpose4_out1);
let linear2_out1 = self.linear2.forward(transpose5_out1);
let unsqueeze1_out1: Tensor<4> = linear1_out1.unsqueeze_dims::<4>(&[2]);
let unsqueeze2_out1: Tensor<4> = linear2_out1.unsqueeze_dims::<4>(&[1]);
let add1_out1 = unsqueeze1_out1.add(unsqueeze2_out1);
let relu1_out1 = burn::tensor::activation::relu(add1_out1);
let linear3_out1 = self.linear3.forward(relu1_out1);
(linear3_out1, target_length, concat2_out1, concat3_out1)
}
}