extern crate alloc;
use burn::prelude::*;
use burn::nn::LayerNorm;
use burn::nn::LayerNormConfig;
use burn::nn::Linear;
use burn::nn::LinearConfig;
#[derive(Module, Debug)]
pub struct Submodule1 {
constant1: burn::module::Param<Tensor<2>>,
constant2: burn::module::Param<Tensor<2>>,
layernormalization1: LayerNorm,
linear1: Linear,
linear2: Linear,
linear3: Linear,
linear4: Linear,
linear5: Linear,
linear6: Linear,
linear7: Linear,
linear8: Linear,
linear9: Linear,
linear10: Linear,
linear11: Linear,
linear12: Linear,
linear13: Linear,
linear14: Linear,
linear15: Linear,
linear16: Linear,
layernormalization2: LayerNorm,
linear17: Linear,
linear18: Linear,
linear19: Linear,
constant42: burn::module::Param<Tensor<1>>,
constant43: burn::module::Param<Tensor<1>>,
constant44: burn::module::Param<Tensor<1>>,
constant237: burn::module::Param<Tensor<1>>,
constant238: burn::module::Param<Tensor<1>>,
linear20: Linear,
#[module(skip)]
device: Device,
}
impl Submodule1 {
#[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([16384, 1024], (device, burn::tensor::DType::F32)),
device.clone(),
false,
[16384, 1024].into(),
);
let constant2: burn::module::Param<Tensor<2>> = burn::module::Param::uninitialized(
burn::module::ParamId::new(),
move |device, _require_grad| Tensor::<
2,
>::zeros([1024, 1024], (device, burn::tensor::DType::F32)),
device.clone(),
false,
[1024, 1024].into(),
);
let layernormalization1 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let linear1 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
let linear2 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
let linear3 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
let linear4 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
let linear5 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
let linear6 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
let linear7 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
let linear8 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
let linear9 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
let linear10 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
let linear11 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
let linear12 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
let linear13 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
let linear14 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
let linear15 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
let linear16 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
let layernormalization2 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let linear17 = LinearConfig::new(1024, 1024).with_bias(false).init(device);
let linear18 = LinearConfig::new(1024, 1024).with_bias(false).init(device);
let linear19 = LinearConfig::new(1024, 1024).with_bias(false).init(device);
let constant42: burn::module::Param<Tensor<1>> = burn::module::Param::uninitialized(
burn::module::ParamId::new(),
move |device, _require_grad| Tensor::<
1,
>::zeros([1024], (device, burn::tensor::DType::F32)),
device.clone(),
false,
[1024].into(),
);
let constant43: burn::module::Param<Tensor<1>> = burn::module::Param::uninitialized(
burn::module::ParamId::new(),
move |device, _require_grad| Tensor::<
1,
>::zeros([1024], (device, burn::tensor::DType::F32)),
device.clone(),
false,
[1024].into(),
);
let constant44: burn::module::Param<Tensor<1>> = burn::module::Param::uninitialized(
burn::module::ParamId::new(),
move |device, _require_grad| Tensor::<
1,
>::zeros([1024], (device, burn::tensor::DType::F32)),
device.clone(),
false,
[1024].into(),
);
let constant237: burn::module::Param<Tensor<1>> = burn::module::Param::uninitialized(
burn::module::ParamId::new(),
move |device, _require_grad| Tensor::<
1,
>::from_data(
burn::tensor::TensorData::from([1f64]),
(device, burn::tensor::DType::F32),
),
device.clone(),
false,
[1].into(),
);
let constant238: burn::module::Param<Tensor<1>> = burn::module::Param::uninitialized(
burn::module::ParamId::new(),
move |device, _require_grad| Tensor::<
1,
>::from_data(
burn::tensor::TensorData::from([
-340282346638528860000000000000000000000f64,
]),
(device, burn::tensor::DType::F32),
),
device.clone(),
false,
[1].into(),
);
let linear20 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
Self {
constant1,
constant2,
layernormalization1,
linear1,
linear2,
linear3,
linear4,
linear5,
linear6,
linear7,
linear8,
linear9,
linear10,
linear11,
linear12,
linear13,
linear14,
linear15,
linear16,
layernormalization2,
linear17,
linear18,
linear19,
constant42,
constant43,
constant44,
constant237,
constant238,
linear20,
device: device.clone(),
}
}
#[allow(clippy::let_and_return, clippy::approx_constant)]
pub fn forward(
&self,
attention_mask: Tensor<2, Int>,
input_ids: Tensor<2, Int>,
position_ids: Tensor<2, Int>,
encoder_hidden_states: Tensor<3>,
) -> (
Tensor<3>,
Tensor<4>,
Tensor<4>,
Tensor<4>,
Tensor<4>,
Tensor<4>,
Tensor<4>,
Tensor<4>,
Tensor<4>,
Tensor<4>,
Tensor<4>,
Tensor<4>,
Tensor<4>,
Tensor<4>,
Tensor<4>,
Tensor<4>,
Tensor<4>,
Tensor<4>,
Tensor<4>,
Tensor<4>,
) {
let shape1_out1: [i64; 2] = {
let axes = &attention_mask.clone().dims()[0..2];
let mut output = [0i64; 2];
for i in 0..2 {
output[i] = axes[i] as i64;
}
output
};
let gather1_out1 = shape1_out1[1] as i64;
let shape2_out1: [i64; 2] = {
let axes = &input_ids.clone().dims()[0..2];
let mut output = [0i64; 2];
for i in 0..2 {
output[i] = axes[i] as i64;
}
output
};
let gather2_out1 = shape2_out1[0] as i64;
let gather3_out1 = shape2_out1[1] as i64;
let unsqueeze1_out1 = [gather2_out1 as i64];
let constant1_out1 = self.constant1.val();
let gather4_out1 = {
let axis_size = constant1_out1.dims()[0] as i64;
let negative = input_ids.clone().lower_elem(0i64);
let corrected = input_ids.clone() + axis_size;
let indices = input_ids.mask_where(negative, corrected);
constant1_out1.take::<2, 3>(0, indices)
};
let constant2_out1 = self.constant2.val();
let gather5_out1 = {
let axis_size = constant2_out1.dims()[0] as i64;
let negative = position_ids.clone().lower_elem(0i64);
let corrected = position_ids.clone() + axis_size;
let indices = position_ids.mask_where(negative, corrected);
constant2_out1.take::<2, 3>(0, indices)
};
let add1_out1 = gather4_out1.add(gather5_out1);
let layernormalization1_out1 = {
let dtype = add1_out1.dtype();
self.layernormalization1
.forward(add1_out1.cast(burn::tensor::DType::F32))
.cast(dtype)
};
let linear1_out1 = self.linear1.forward(encoder_hidden_states.clone());
let reshape1_out1 = linear1_out1.reshape([0, -1, 8, 128]);
let transpose1_out1 = reshape1_out1.permute([0, 2, 1, 3]);
let linear2_out1 = self.linear2.forward(encoder_hidden_states.clone());
let reshape2_out1 = linear2_out1.reshape([0, -1, 8, 128]);
let transpose2_out1 = reshape2_out1.permute([0, 2, 1, 3]);
let linear3_out1 = self.linear3.forward(encoder_hidden_states.clone());
let reshape3_out1 = linear3_out1.reshape([0, -1, 8, 128]);
let transpose3_out1 = reshape3_out1.permute([0, 2, 1, 3]);
let linear4_out1 = self.linear4.forward(encoder_hidden_states.clone());
let reshape4_out1 = linear4_out1.reshape([0, -1, 8, 128]);
let transpose4_out1 = reshape4_out1.permute([0, 2, 1, 3]);
let linear5_out1 = self.linear5.forward(encoder_hidden_states.clone());
let reshape5_out1 = linear5_out1.reshape([0, -1, 8, 128]);
let transpose5_out1 = reshape5_out1.permute([0, 2, 1, 3]);
let linear6_out1 = self.linear6.forward(encoder_hidden_states.clone());
let reshape6_out1 = linear6_out1.reshape([0, -1, 8, 128]);
let transpose6_out1 = reshape6_out1.permute([0, 2, 1, 3]);
let linear7_out1 = self.linear7.forward(encoder_hidden_states.clone());
let reshape7_out1 = linear7_out1.reshape([0, -1, 8, 128]);
let transpose7_out1 = reshape7_out1.permute([0, 2, 1, 3]);
let linear8_out1 = self.linear8.forward(encoder_hidden_states.clone());
let reshape8_out1 = linear8_out1.reshape([0, -1, 8, 128]);
let transpose8_out1 = reshape8_out1.permute([0, 2, 1, 3]);
let linear9_out1 = self.linear9.forward(encoder_hidden_states.clone());
let reshape9_out1 = linear9_out1.reshape([0, -1, 8, 128]);
let transpose9_out1 = reshape9_out1.permute([0, 2, 1, 3]);
let linear10_out1 = self.linear10.forward(encoder_hidden_states.clone());
let reshape10_out1 = linear10_out1.reshape([0, -1, 8, 128]);
let transpose10_out1 = reshape10_out1.permute([0, 2, 1, 3]);
let linear11_out1 = self.linear11.forward(encoder_hidden_states.clone());
let reshape11_out1 = linear11_out1.reshape([0, -1, 8, 128]);
let transpose11_out1 = reshape11_out1.permute([0, 2, 1, 3]);
let linear12_out1 = self.linear12.forward(encoder_hidden_states.clone());
let reshape12_out1 = linear12_out1.reshape([0, -1, 8, 128]);
let transpose12_out1 = reshape12_out1.permute([0, 2, 1, 3]);
let linear13_out1 = self.linear13.forward(encoder_hidden_states.clone());
let reshape13_out1 = linear13_out1.reshape([0, -1, 8, 128]);
let transpose13_out1 = reshape13_out1.permute([0, 2, 1, 3]);
let linear14_out1 = self.linear14.forward(encoder_hidden_states.clone());
let reshape14_out1 = linear14_out1.reshape([0, -1, 8, 128]);
let transpose14_out1 = reshape14_out1.permute([0, 2, 1, 3]);
let linear15_out1 = self.linear15.forward(encoder_hidden_states.clone());
let reshape15_out1 = linear15_out1.reshape([0, -1, 8, 128]);
let transpose15_out1 = reshape15_out1.permute([0, 2, 1, 3]);
let linear16_out1 = self.linear16.forward(encoder_hidden_states);
let reshape16_out1 = linear16_out1.reshape([0, -1, 8, 128]);
let transpose16_out1 = reshape16_out1.permute([0, 2, 1, 3]);
let layernormalization2_out1 = {
let dtype = layernormalization1_out1.clone().dtype();
self.layernormalization2
.forward(layernormalization1_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let linear17_out1 = self.linear17.forward(layernormalization2_out1.clone());
let linear18_out1 = self.linear18.forward(layernormalization2_out1.clone());
let linear19_out1 = self.linear19.forward(layernormalization2_out1);
let constant42_out1 = self.constant42.val();
let add2_out1 = linear17_out1
.add((constant42_out1).unsqueeze_dims(&[0isize, 1isize]));
let constant43_out1 = self.constant43.val();
let add3_out1 = linear18_out1
.add((constant43_out1).unsqueeze_dims(&[0isize, 1isize]));
let constant44_out1 = self.constant44.val();
let add4_out1 = linear19_out1
.add((constant44_out1).unsqueeze_dims(&[0isize, 1isize]));
let cast1_out1 = attention_mask.float().cast(burn::tensor::DType::F32);
let constant237_out1 = self.constant237.val();
let sub1_out1 = (constant237_out1).unsqueeze_dims(&[0isize]).sub(cast1_out1);
let constant238_out1 = self.constant238.val();
let mul1_out1 = sub1_out1.mul((constant238_out1).unsqueeze_dims(&[0isize]));
let unsqueeze2_out1: Tensor<4> = mul1_out1.unsqueeze_dims::<4>(&[1, 2]);
let unsqueeze3_out1 = [gather3_out1 as i64];
let unsqueeze4_out1 = [gather1_out1 as i64];
let constant242_out1: [i64; 1] = [1i64];
let concat1_out1: [i64; 4usize] = [
&unsqueeze1_out1[..],
&constant242_out1[..],
&unsqueeze3_out1[..],
&unsqueeze4_out1[..],
]
.concat()
.try_into()
.unwrap();
let expand1_out1 = {
let onnx_shape: [i64; 4usize] = concat1_out1;
let input_dims = unsqueeze2_out1.dims();
let mut shape = onnx_shape;
#[allow(clippy::needless_range_loop)]
for i in 0..4usize {
let dim_offset = i;
if shape[dim_offset] == 1 && input_dims[i] > 1 {
shape[dim_offset] = input_dims[i] as i64;
}
}
unsqueeze2_out1.expand(shape)
};
let reshape17_out1 = add3_out1.reshape([0, 0, 8, -1]);
let transpose17_out1 = reshape17_out1.permute([0, 2, 1, 3]);
let reshape18_out1 = add4_out1.reshape([0, 0, 8, -1]);
let transpose18_out1 = reshape18_out1.permute([0, 2, 1, 3]);
let reshape19_out1 = add2_out1.reshape([0, 0, 8, -1]);
let transpose19_out1 = reshape19_out1.permute([0, 2, 1, 3]);
let (attention1_out1,) = {
let q = transpose19_out1;
let k = transpose17_out1.clone();
let v = transpose18_out1.clone();
let hide = {
let [batch, heads, q_len, _] = q.dims();
let k_len = k.dims()[2];
let rows = Tensor::<
1,
Int,
>::arange(0..q_len as i64, (&self.device, burn::tensor::DType::I64))
.reshape([q_len, 1])
.expand([q_len, k_len]);
let cols = Tensor::<
1,
Int,
>::arange(0..k_len as i64, (&self.device, burn::tensor::DType::I64))
.reshape([1, k_len])
.expand([q_len, k_len]);
cols.greater(rows).unsqueeze::<4>().expand([batch, heads, q_len, k_len])
};
let score_bias = expand1_out1.clone();
let hidden_rows = {
let shape = [q.dims()[0], q.dims()[1], q.dims()[2], k.dims()[2]];
score_bias
.clone()
.expand(shape)
.mask_fill(hide.clone().expand(shape), f32::NEG_INFINITY)
.max_dim(3)
.equal_elem(f32::NEG_INFINITY)
};
let attention1_out1 = burn::tensor::module::attention(
q,
k,
v,
Some(hide),
Some(score_bias),
burn::tensor::ops::AttentionModuleOptions {
scale: Some(0.0883883461356163f64),
softcap: None,
is_causal: false,
},
);
let attention1_out1 = {
let dims = attention1_out1.dims();
attention1_out1.mask_fill(hidden_rows.expand(dims), 0.0)
};
(attention1_out1,)
};
let transpose20_out1 = attention1_out1.permute([0, 2, 1, 3]);
let reshape20_out1 = transpose20_out1.reshape([0, 0, -1]);
let linear20_out1 = self.linear20.forward(reshape20_out1);
let add5_out1 = layernormalization1_out1.add(linear20_out1);
(
add5_out1,
transpose1_out1,
transpose2_out1,
expand1_out1,
transpose3_out1,
transpose4_out1,
transpose5_out1,
transpose6_out1,
transpose7_out1,
transpose8_out1,
transpose9_out1,
transpose10_out1,
transpose11_out1,
transpose12_out1,
transpose13_out1,
transpose14_out1,
transpose15_out1,
transpose16_out1,
transpose17_out1,
transpose18_out1,
)
}
}
#[derive(Module, Debug)]
pub struct Submodule2 {
layernormalization3: LayerNorm,
linear21: Linear,
linear22: Linear,
layernormalization4: LayerNorm,
linear23: Linear,
linear24: Linear,
layernormalization5: LayerNorm,
linear25: Linear,
linear26: Linear,
linear27: Linear,
constant64: burn::module::Param<Tensor<1>>,
constant65: burn::module::Param<Tensor<1>>,
constant66: burn::module::Param<Tensor<1>>,
linear28: Linear,
layernormalization6: LayerNorm,
linear29: Linear,
linear30: Linear,
layernormalization7: LayerNorm,
linear31: Linear,
linear32: Linear,
layernormalization8: LayerNorm,
linear33: Linear,
linear34: Linear,
linear35: Linear,
constant86: burn::module::Param<Tensor<1>>,
constant87: burn::module::Param<Tensor<1>>,
constant88: burn::module::Param<Tensor<1>>,
linear36: Linear,
#[module(skip)]
device: Device,
}
impl Submodule2 {
#[allow(unused_variables)]
pub fn new(device: &Device) -> Self {
let layernormalization3 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let linear21 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
let linear22 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
let layernormalization4 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let linear23 = LinearConfig::new(1024, 4096).with_bias(true).init(device);
let linear24 = LinearConfig::new(4096, 1024).with_bias(true).init(device);
let layernormalization5 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let linear25 = LinearConfig::new(1024, 1024).with_bias(false).init(device);
let linear26 = LinearConfig::new(1024, 1024).with_bias(false).init(device);
let linear27 = LinearConfig::new(1024, 1024).with_bias(false).init(device);
let constant64: burn::module::Param<Tensor<1>> = burn::module::Param::uninitialized(
burn::module::ParamId::new(),
move |device, _require_grad| Tensor::<
1,
>::zeros([1024], (device, burn::tensor::DType::F32)),
device.clone(),
false,
[1024].into(),
);
let constant65: burn::module::Param<Tensor<1>> = burn::module::Param::uninitialized(
burn::module::ParamId::new(),
move |device, _require_grad| Tensor::<
1,
>::zeros([1024], (device, burn::tensor::DType::F32)),
device.clone(),
false,
[1024].into(),
);
let constant66: burn::module::Param<Tensor<1>> = burn::module::Param::uninitialized(
burn::module::ParamId::new(),
move |device, _require_grad| Tensor::<
1,
>::zeros([1024], (device, burn::tensor::DType::F32)),
device.clone(),
false,
[1024].into(),
);
let linear28 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
let layernormalization6 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let linear29 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
let linear30 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
let layernormalization7 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let linear31 = LinearConfig::new(1024, 4096).with_bias(true).init(device);
let linear32 = LinearConfig::new(4096, 1024).with_bias(true).init(device);
let layernormalization8 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let linear33 = LinearConfig::new(1024, 1024).with_bias(false).init(device);
let linear34 = LinearConfig::new(1024, 1024).with_bias(false).init(device);
let linear35 = LinearConfig::new(1024, 1024).with_bias(false).init(device);
let constant86: burn::module::Param<Tensor<1>> = burn::module::Param::uninitialized(
burn::module::ParamId::new(),
move |device, _require_grad| Tensor::<
1,
>::zeros([1024], (device, burn::tensor::DType::F32)),
device.clone(),
false,
[1024].into(),
);
let constant87: burn::module::Param<Tensor<1>> = burn::module::Param::uninitialized(
burn::module::ParamId::new(),
move |device, _require_grad| Tensor::<
1,
>::zeros([1024], (device, burn::tensor::DType::F32)),
device.clone(),
false,
[1024].into(),
);
let constant88: burn::module::Param<Tensor<1>> = burn::module::Param::uninitialized(
burn::module::ParamId::new(),
move |device, _require_grad| Tensor::<
1,
>::zeros([1024], (device, burn::tensor::DType::F32)),
device.clone(),
false,
[1024].into(),
);
let linear36 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
Self {
layernormalization3,
linear21,
linear22,
layernormalization4,
linear23,
linear24,
layernormalization5,
linear25,
linear26,
linear27,
constant64,
constant65,
constant66,
linear28,
layernormalization6,
linear29,
linear30,
layernormalization7,
linear31,
linear32,
layernormalization8,
linear33,
linear34,
linear35,
constant86,
constant87,
constant88,
linear36,
device: device.clone(),
}
}
#[allow(clippy::let_and_return, clippy::approx_constant)]
pub fn forward(
&self,
add5_out1: Tensor<3>,
transpose1_out1: Tensor<4>,
transpose2_out1: Tensor<4>,
expand1_out1: Tensor<4>,
transpose3_out1: Tensor<4>,
transpose4_out1: Tensor<4>,
) -> (Tensor<3>, Tensor<4>, Tensor<4>, Tensor<4>, Tensor<4>) {
let layernormalization3_out1 = {
let dtype = add5_out1.clone().dtype();
self.layernormalization3
.forward(add5_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let linear21_out1 = self.linear21.forward(layernormalization3_out1);
let reshape21_out1 = linear21_out1.reshape([0, 0, 8, -1]);
let transpose21_out1 = reshape21_out1.permute([0, 2, 1, 3]);
let (attention2_out1,) = {
let q = transpose21_out1;
let k = transpose1_out1;
let v = transpose2_out1;
let attention2_out1 = burn::tensor::module::attention(
q,
k,
v,
None,
None,
burn::tensor::ops::AttentionModuleOptions {
scale: Some(0.0883883461356163f64),
softcap: None,
is_causal: false,
},
);
(attention2_out1,)
};
let transpose22_out1 = attention2_out1.permute([0, 2, 1, 3]);
let reshape22_out1 = transpose22_out1.reshape([0, 0, -1]);
let linear22_out1 = self.linear22.forward(reshape22_out1);
let add6_out1 = add5_out1.add(linear22_out1);
let layernormalization4_out1 = {
let dtype = add6_out1.clone().dtype();
self.layernormalization4
.forward(add6_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let linear23_out1 = self.linear23.forward(layernormalization4_out1);
let relu1_out1 = burn::tensor::activation::relu(linear23_out1);
let linear24_out1 = self.linear24.forward(relu1_out1);
let add7_out1 = add6_out1.add(linear24_out1);
let layernormalization5_out1 = {
let dtype = add7_out1.clone().dtype();
self.layernormalization5
.forward(add7_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let linear25_out1 = self.linear25.forward(layernormalization5_out1.clone());
let linear26_out1 = self.linear26.forward(layernormalization5_out1.clone());
let linear27_out1 = self.linear27.forward(layernormalization5_out1);
let constant64_out1 = self.constant64.val();
let add8_out1 = linear25_out1
.add((constant64_out1).unsqueeze_dims(&[0isize, 1isize]));
let constant65_out1 = self.constant65.val();
let add9_out1 = linear26_out1
.add((constant65_out1).unsqueeze_dims(&[0isize, 1isize]));
let constant66_out1 = self.constant66.val();
let add10_out1 = linear27_out1
.add((constant66_out1).unsqueeze_dims(&[0isize, 1isize]));
let reshape23_out1 = add9_out1.reshape([0, 0, 8, -1]);
let transpose23_out1 = reshape23_out1.permute([0, 2, 1, 3]);
let reshape24_out1 = add10_out1.reshape([0, 0, 8, -1]);
let transpose24_out1 = reshape24_out1.permute([0, 2, 1, 3]);
let reshape25_out1 = add8_out1.reshape([0, 0, 8, -1]);
let transpose25_out1 = reshape25_out1.permute([0, 2, 1, 3]);
let (attention3_out1,) = {
let q = transpose25_out1;
let k = transpose23_out1.clone();
let v = transpose24_out1.clone();
let hide = {
let [batch, heads, q_len, _] = q.dims();
let k_len = k.dims()[2];
let rows = Tensor::<
1,
Int,
>::arange(0..q_len as i64, (&self.device, burn::tensor::DType::I64))
.reshape([q_len, 1])
.expand([q_len, k_len]);
let cols = Tensor::<
1,
Int,
>::arange(0..k_len as i64, (&self.device, burn::tensor::DType::I64))
.reshape([1, k_len])
.expand([q_len, k_len]);
cols.greater(rows).unsqueeze::<4>().expand([batch, heads, q_len, k_len])
};
let score_bias = expand1_out1.clone();
let hidden_rows = {
let shape = [q.dims()[0], q.dims()[1], q.dims()[2], k.dims()[2]];
score_bias
.clone()
.expand(shape)
.mask_fill(hide.clone().expand(shape), f32::NEG_INFINITY)
.max_dim(3)
.equal_elem(f32::NEG_INFINITY)
};
let attention3_out1 = burn::tensor::module::attention(
q,
k,
v,
Some(hide),
Some(score_bias),
burn::tensor::ops::AttentionModuleOptions {
scale: Some(0.0883883461356163f64),
softcap: None,
is_causal: false,
},
);
let attention3_out1 = {
let dims = attention3_out1.dims();
attention3_out1.mask_fill(hidden_rows.expand(dims), 0.0)
};
(attention3_out1,)
};
let transpose26_out1 = attention3_out1.permute([0, 2, 1, 3]);
let reshape26_out1 = transpose26_out1.reshape([0, 0, -1]);
let linear28_out1 = self.linear28.forward(reshape26_out1);
let add11_out1 = add7_out1.add(linear28_out1);
let layernormalization6_out1 = {
let dtype = add11_out1.clone().dtype();
self.layernormalization6
.forward(add11_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let linear29_out1 = self.linear29.forward(layernormalization6_out1);
let reshape27_out1 = linear29_out1.reshape([0, 0, 8, -1]);
let transpose27_out1 = reshape27_out1.permute([0, 2, 1, 3]);
let (attention4_out1,) = {
let q = transpose27_out1;
let k = transpose3_out1;
let v = transpose4_out1;
let attention4_out1 = burn::tensor::module::attention(
q,
k,
v,
None,
None,
burn::tensor::ops::AttentionModuleOptions {
scale: Some(0.0883883461356163f64),
softcap: None,
is_causal: false,
},
);
(attention4_out1,)
};
let transpose28_out1 = attention4_out1.permute([0, 2, 1, 3]);
let reshape28_out1 = transpose28_out1.reshape([0, 0, -1]);
let linear30_out1 = self.linear30.forward(reshape28_out1);
let add12_out1 = add11_out1.add(linear30_out1);
let layernormalization7_out1 = {
let dtype = add12_out1.clone().dtype();
self.layernormalization7
.forward(add12_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let linear31_out1 = self.linear31.forward(layernormalization7_out1);
let relu2_out1 = burn::tensor::activation::relu(linear31_out1);
let linear32_out1 = self.linear32.forward(relu2_out1);
let add13_out1 = add12_out1.add(linear32_out1);
let layernormalization8_out1 = {
let dtype = add13_out1.clone().dtype();
self.layernormalization8
.forward(add13_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let linear33_out1 = self.linear33.forward(layernormalization8_out1.clone());
let linear34_out1 = self.linear34.forward(layernormalization8_out1.clone());
let linear35_out1 = self.linear35.forward(layernormalization8_out1);
let constant86_out1 = self.constant86.val();
let add14_out1 = linear33_out1
.add((constant86_out1).unsqueeze_dims(&[0isize, 1isize]));
let constant87_out1 = self.constant87.val();
let add15_out1 = linear34_out1
.add((constant87_out1).unsqueeze_dims(&[0isize, 1isize]));
let constant88_out1 = self.constant88.val();
let add16_out1 = linear35_out1
.add((constant88_out1).unsqueeze_dims(&[0isize, 1isize]));
let reshape29_out1 = add15_out1.reshape([0, 0, 8, -1]);
let transpose29_out1 = reshape29_out1.permute([0, 2, 1, 3]);
let reshape30_out1 = add16_out1.reshape([0, 0, 8, -1]);
let transpose30_out1 = reshape30_out1.permute([0, 2, 1, 3]);
let reshape31_out1 = add14_out1.reshape([0, 0, 8, -1]);
let transpose31_out1 = reshape31_out1.permute([0, 2, 1, 3]);
let (attention5_out1,) = {
let q = transpose31_out1;
let k = transpose29_out1.clone();
let v = transpose30_out1.clone();
let hide = {
let [batch, heads, q_len, _] = q.dims();
let k_len = k.dims()[2];
let rows = Tensor::<
1,
Int,
>::arange(0..q_len as i64, (&self.device, burn::tensor::DType::I64))
.reshape([q_len, 1])
.expand([q_len, k_len]);
let cols = Tensor::<
1,
Int,
>::arange(0..k_len as i64, (&self.device, burn::tensor::DType::I64))
.reshape([1, k_len])
.expand([q_len, k_len]);
cols.greater(rows).unsqueeze::<4>().expand([batch, heads, q_len, k_len])
};
let score_bias = expand1_out1;
let hidden_rows = {
let shape = [q.dims()[0], q.dims()[1], q.dims()[2], k.dims()[2]];
score_bias
.clone()
.expand(shape)
.mask_fill(hide.clone().expand(shape), f32::NEG_INFINITY)
.max_dim(3)
.equal_elem(f32::NEG_INFINITY)
};
let attention5_out1 = burn::tensor::module::attention(
q,
k,
v,
Some(hide),
Some(score_bias),
burn::tensor::ops::AttentionModuleOptions {
scale: Some(0.0883883461356163f64),
softcap: None,
is_causal: false,
},
);
let attention5_out1 = {
let dims = attention5_out1.dims();
attention5_out1.mask_fill(hidden_rows.expand(dims), 0.0)
};
(attention5_out1,)
};
let transpose32_out1 = attention5_out1.permute([0, 2, 1, 3]);
let reshape32_out1 = transpose32_out1.reshape([0, 0, -1]);
let linear36_out1 = self.linear36.forward(reshape32_out1);
let add17_out1 = add13_out1.add(linear36_out1);
(
add17_out1,
transpose23_out1,
transpose24_out1,
transpose29_out1,
transpose30_out1,
)
}
}
#[derive(Module, Debug)]
pub struct Submodule3 {
layernormalization9: LayerNorm,
linear37: Linear,
linear38: Linear,
layernormalization10: LayerNorm,
linear39: Linear,
linear40: Linear,
layernormalization11: LayerNorm,
linear41: Linear,
linear42: Linear,
linear43: Linear,
constant108: burn::module::Param<Tensor<1>>,
constant109: burn::module::Param<Tensor<1>>,
constant110: burn::module::Param<Tensor<1>>,
linear44: Linear,
layernormalization12: LayerNorm,
linear45: Linear,
linear46: Linear,
layernormalization13: LayerNorm,
linear47: Linear,
linear48: Linear,
layernormalization14: LayerNorm,
linear49: Linear,
linear50: Linear,
linear51: Linear,
constant130: burn::module::Param<Tensor<1>>,
constant131: burn::module::Param<Tensor<1>>,
constant132: burn::module::Param<Tensor<1>>,
linear52: Linear,
#[module(skip)]
device: Device,
}
impl Submodule3 {
#[allow(unused_variables)]
pub fn new(device: &Device) -> Self {
let layernormalization9 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let linear37 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
let linear38 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
let layernormalization10 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let linear39 = LinearConfig::new(1024, 4096).with_bias(true).init(device);
let linear40 = LinearConfig::new(4096, 1024).with_bias(true).init(device);
let layernormalization11 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let linear41 = LinearConfig::new(1024, 1024).with_bias(false).init(device);
let linear42 = LinearConfig::new(1024, 1024).with_bias(false).init(device);
let linear43 = LinearConfig::new(1024, 1024).with_bias(false).init(device);
let constant108: burn::module::Param<Tensor<1>> = burn::module::Param::uninitialized(
burn::module::ParamId::new(),
move |device, _require_grad| Tensor::<
1,
>::zeros([1024], (device, burn::tensor::DType::F32)),
device.clone(),
false,
[1024].into(),
);
let constant109: burn::module::Param<Tensor<1>> = burn::module::Param::uninitialized(
burn::module::ParamId::new(),
move |device, _require_grad| Tensor::<
1,
>::zeros([1024], (device, burn::tensor::DType::F32)),
device.clone(),
false,
[1024].into(),
);
let constant110: burn::module::Param<Tensor<1>> = burn::module::Param::uninitialized(
burn::module::ParamId::new(),
move |device, _require_grad| Tensor::<
1,
>::zeros([1024], (device, burn::tensor::DType::F32)),
device.clone(),
false,
[1024].into(),
);
let linear44 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
let layernormalization12 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let linear45 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
let linear46 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
let layernormalization13 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let linear47 = LinearConfig::new(1024, 4096).with_bias(true).init(device);
let linear48 = LinearConfig::new(4096, 1024).with_bias(true).init(device);
let layernormalization14 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let linear49 = LinearConfig::new(1024, 1024).with_bias(false).init(device);
let linear50 = LinearConfig::new(1024, 1024).with_bias(false).init(device);
let linear51 = LinearConfig::new(1024, 1024).with_bias(false).init(device);
let constant130: burn::module::Param<Tensor<1>> = burn::module::Param::uninitialized(
burn::module::ParamId::new(),
move |device, _require_grad| Tensor::<
1,
>::zeros([1024], (device, burn::tensor::DType::F32)),
device.clone(),
false,
[1024].into(),
);
let constant131: burn::module::Param<Tensor<1>> = burn::module::Param::uninitialized(
burn::module::ParamId::new(),
move |device, _require_grad| Tensor::<
1,
>::zeros([1024], (device, burn::tensor::DType::F32)),
device.clone(),
false,
[1024].into(),
);
let constant132: burn::module::Param<Tensor<1>> = burn::module::Param::uninitialized(
burn::module::ParamId::new(),
move |device, _require_grad| Tensor::<
1,
>::zeros([1024], (device, burn::tensor::DType::F32)),
device.clone(),
false,
[1024].into(),
);
let linear52 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
Self {
layernormalization9,
linear37,
linear38,
layernormalization10,
linear39,
linear40,
layernormalization11,
linear41,
linear42,
linear43,
constant108,
constant109,
constant110,
linear44,
layernormalization12,
linear45,
linear46,
layernormalization13,
linear47,
linear48,
layernormalization14,
linear49,
linear50,
linear51,
constant130,
constant131,
constant132,
linear52,
device: device.clone(),
}
}
#[allow(clippy::let_and_return, clippy::approx_constant)]
pub fn forward(
&self,
add17_out1: Tensor<3>,
transpose5_out1: Tensor<4>,
transpose6_out1: Tensor<4>,
expand1_out1: Tensor<4>,
transpose7_out1: Tensor<4>,
transpose8_out1: Tensor<4>,
) -> (Tensor<3>, Tensor<4>, Tensor<4>, Tensor<4>, Tensor<4>) {
let layernormalization9_out1 = {
let dtype = add17_out1.clone().dtype();
self.layernormalization9
.forward(add17_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let linear37_out1 = self.linear37.forward(layernormalization9_out1);
let reshape33_out1 = linear37_out1.reshape([0, 0, 8, -1]);
let transpose33_out1 = reshape33_out1.permute([0, 2, 1, 3]);
let (attention6_out1,) = {
let q = transpose33_out1;
let k = transpose5_out1;
let v = transpose6_out1;
let attention6_out1 = burn::tensor::module::attention(
q,
k,
v,
None,
None,
burn::tensor::ops::AttentionModuleOptions {
scale: Some(0.0883883461356163f64),
softcap: None,
is_causal: false,
},
);
(attention6_out1,)
};
let transpose34_out1 = attention6_out1.permute([0, 2, 1, 3]);
let reshape34_out1 = transpose34_out1.reshape([0, 0, -1]);
let linear38_out1 = self.linear38.forward(reshape34_out1);
let add18_out1 = add17_out1.add(linear38_out1);
let layernormalization10_out1 = {
let dtype = add18_out1.clone().dtype();
self.layernormalization10
.forward(add18_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let linear39_out1 = self.linear39.forward(layernormalization10_out1);
let relu3_out1 = burn::tensor::activation::relu(linear39_out1);
let linear40_out1 = self.linear40.forward(relu3_out1);
let add19_out1 = add18_out1.add(linear40_out1);
let layernormalization11_out1 = {
let dtype = add19_out1.clone().dtype();
self.layernormalization11
.forward(add19_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let linear41_out1 = self.linear41.forward(layernormalization11_out1.clone());
let linear42_out1 = self.linear42.forward(layernormalization11_out1.clone());
let linear43_out1 = self.linear43.forward(layernormalization11_out1);
let constant108_out1 = self.constant108.val();
let add20_out1 = linear41_out1
.add((constant108_out1).unsqueeze_dims(&[0isize, 1isize]));
let constant109_out1 = self.constant109.val();
let add21_out1 = linear42_out1
.add((constant109_out1).unsqueeze_dims(&[0isize, 1isize]));
let constant110_out1 = self.constant110.val();
let add22_out1 = linear43_out1
.add((constant110_out1).unsqueeze_dims(&[0isize, 1isize]));
let reshape35_out1 = add21_out1.reshape([0, 0, 8, -1]);
let transpose35_out1 = reshape35_out1.permute([0, 2, 1, 3]);
let reshape36_out1 = add22_out1.reshape([0, 0, 8, -1]);
let transpose36_out1 = reshape36_out1.permute([0, 2, 1, 3]);
let reshape37_out1 = add20_out1.reshape([0, 0, 8, -1]);
let transpose37_out1 = reshape37_out1.permute([0, 2, 1, 3]);
let (attention7_out1,) = {
let q = transpose37_out1;
let k = transpose35_out1.clone();
let v = transpose36_out1.clone();
let hide = {
let [batch, heads, q_len, _] = q.dims();
let k_len = k.dims()[2];
let rows = Tensor::<
1,
Int,
>::arange(0..q_len as i64, (&self.device, burn::tensor::DType::I64))
.reshape([q_len, 1])
.expand([q_len, k_len]);
let cols = Tensor::<
1,
Int,
>::arange(0..k_len as i64, (&self.device, burn::tensor::DType::I64))
.reshape([1, k_len])
.expand([q_len, k_len]);
cols.greater(rows).unsqueeze::<4>().expand([batch, heads, q_len, k_len])
};
let score_bias = expand1_out1.clone();
let hidden_rows = {
let shape = [q.dims()[0], q.dims()[1], q.dims()[2], k.dims()[2]];
score_bias
.clone()
.expand(shape)
.mask_fill(hide.clone().expand(shape), f32::NEG_INFINITY)
.max_dim(3)
.equal_elem(f32::NEG_INFINITY)
};
let attention7_out1 = burn::tensor::module::attention(
q,
k,
v,
Some(hide),
Some(score_bias),
burn::tensor::ops::AttentionModuleOptions {
scale: Some(0.0883883461356163f64),
softcap: None,
is_causal: false,
},
);
let attention7_out1 = {
let dims = attention7_out1.dims();
attention7_out1.mask_fill(hidden_rows.expand(dims), 0.0)
};
(attention7_out1,)
};
let transpose38_out1 = attention7_out1.permute([0, 2, 1, 3]);
let reshape38_out1 = transpose38_out1.reshape([0, 0, -1]);
let linear44_out1 = self.linear44.forward(reshape38_out1);
let add23_out1 = add19_out1.add(linear44_out1);
let layernormalization12_out1 = {
let dtype = add23_out1.clone().dtype();
self.layernormalization12
.forward(add23_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let linear45_out1 = self.linear45.forward(layernormalization12_out1);
let reshape39_out1 = linear45_out1.reshape([0, 0, 8, -1]);
let transpose39_out1 = reshape39_out1.permute([0, 2, 1, 3]);
let (attention8_out1,) = {
let q = transpose39_out1;
let k = transpose7_out1;
let v = transpose8_out1;
let attention8_out1 = burn::tensor::module::attention(
q,
k,
v,
None,
None,
burn::tensor::ops::AttentionModuleOptions {
scale: Some(0.0883883461356163f64),
softcap: None,
is_causal: false,
},
);
(attention8_out1,)
};
let transpose40_out1 = attention8_out1.permute([0, 2, 1, 3]);
let reshape40_out1 = transpose40_out1.reshape([0, 0, -1]);
let linear46_out1 = self.linear46.forward(reshape40_out1);
let add24_out1 = add23_out1.add(linear46_out1);
let layernormalization13_out1 = {
let dtype = add24_out1.clone().dtype();
self.layernormalization13
.forward(add24_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let linear47_out1 = self.linear47.forward(layernormalization13_out1);
let relu4_out1 = burn::tensor::activation::relu(linear47_out1);
let linear48_out1 = self.linear48.forward(relu4_out1);
let add25_out1 = add24_out1.add(linear48_out1);
let layernormalization14_out1 = {
let dtype = add25_out1.clone().dtype();
self.layernormalization14
.forward(add25_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let linear49_out1 = self.linear49.forward(layernormalization14_out1.clone());
let linear50_out1 = self.linear50.forward(layernormalization14_out1.clone());
let linear51_out1 = self.linear51.forward(layernormalization14_out1);
let constant130_out1 = self.constant130.val();
let add26_out1 = linear49_out1
.add((constant130_out1).unsqueeze_dims(&[0isize, 1isize]));
let constant131_out1 = self.constant131.val();
let add27_out1 = linear50_out1
.add((constant131_out1).unsqueeze_dims(&[0isize, 1isize]));
let constant132_out1 = self.constant132.val();
let add28_out1 = linear51_out1
.add((constant132_out1).unsqueeze_dims(&[0isize, 1isize]));
let reshape41_out1 = add27_out1.reshape([0, 0, 8, -1]);
let transpose41_out1 = reshape41_out1.permute([0, 2, 1, 3]);
let reshape42_out1 = add28_out1.reshape([0, 0, 8, -1]);
let transpose42_out1 = reshape42_out1.permute([0, 2, 1, 3]);
let reshape43_out1 = add26_out1.reshape([0, 0, 8, -1]);
let transpose43_out1 = reshape43_out1.permute([0, 2, 1, 3]);
let (attention9_out1,) = {
let q = transpose43_out1;
let k = transpose41_out1.clone();
let v = transpose42_out1.clone();
let hide = {
let [batch, heads, q_len, _] = q.dims();
let k_len = k.dims()[2];
let rows = Tensor::<
1,
Int,
>::arange(0..q_len as i64, (&self.device, burn::tensor::DType::I64))
.reshape([q_len, 1])
.expand([q_len, k_len]);
let cols = Tensor::<
1,
Int,
>::arange(0..k_len as i64, (&self.device, burn::tensor::DType::I64))
.reshape([1, k_len])
.expand([q_len, k_len]);
cols.greater(rows).unsqueeze::<4>().expand([batch, heads, q_len, k_len])
};
let score_bias = expand1_out1;
let hidden_rows = {
let shape = [q.dims()[0], q.dims()[1], q.dims()[2], k.dims()[2]];
score_bias
.clone()
.expand(shape)
.mask_fill(hide.clone().expand(shape), f32::NEG_INFINITY)
.max_dim(3)
.equal_elem(f32::NEG_INFINITY)
};
let attention9_out1 = burn::tensor::module::attention(
q,
k,
v,
Some(hide),
Some(score_bias),
burn::tensor::ops::AttentionModuleOptions {
scale: Some(0.0883883461356163f64),
softcap: None,
is_causal: false,
},
);
let attention9_out1 = {
let dims = attention9_out1.dims();
attention9_out1.mask_fill(hidden_rows.expand(dims), 0.0)
};
(attention9_out1,)
};
let transpose44_out1 = attention9_out1.permute([0, 2, 1, 3]);
let reshape44_out1 = transpose44_out1.reshape([0, 0, -1]);
let linear52_out1 = self.linear52.forward(reshape44_out1);
let add29_out1 = add25_out1.add(linear52_out1);
(
add29_out1,
transpose35_out1,
transpose36_out1,
transpose41_out1,
transpose42_out1,
)
}
}
#[derive(Module, Debug)]
pub struct Submodule4 {
layernormalization15: LayerNorm,
linear53: Linear,
linear54: Linear,
layernormalization16: LayerNorm,
linear55: Linear,
linear56: Linear,
layernormalization17: LayerNorm,
linear57: Linear,
linear58: Linear,
linear59: Linear,
constant152: burn::module::Param<Tensor<1>>,
constant153: burn::module::Param<Tensor<1>>,
constant154: burn::module::Param<Tensor<1>>,
linear60: Linear,
layernormalization18: LayerNorm,
linear61: Linear,
linear62: Linear,
layernormalization19: LayerNorm,
linear63: Linear,
linear64: Linear,
layernormalization20: LayerNorm,
linear65: Linear,
linear66: Linear,
linear67: Linear,
constant174: burn::module::Param<Tensor<1>>,
constant175: burn::module::Param<Tensor<1>>,
constant176: burn::module::Param<Tensor<1>>,
linear68: Linear,
layernormalization21: LayerNorm,
linear69: Linear,
linear70: Linear,
layernormalization22: LayerNorm,
linear71: Linear,
linear72: Linear,
layernormalization23: LayerNorm,
linear73: Linear,
linear74: Linear,
linear75: Linear,
constant196: burn::module::Param<Tensor<1>>,
constant197: burn::module::Param<Tensor<1>>,
constant198: burn::module::Param<Tensor<1>>,
linear76: Linear,
layernormalization24: LayerNorm,
linear77: Linear,
linear78: Linear,
layernormalization25: LayerNorm,
linear79: Linear,
linear80: Linear,
layernormalization26: LayerNorm,
linear81: Linear,
#[module(skip)]
device: Device,
}
impl Submodule4 {
#[allow(unused_variables)]
pub fn new(device: &Device) -> Self {
let layernormalization15 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let linear53 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
let linear54 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
let layernormalization16 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let linear55 = LinearConfig::new(1024, 4096).with_bias(true).init(device);
let linear56 = LinearConfig::new(4096, 1024).with_bias(true).init(device);
let layernormalization17 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let linear57 = LinearConfig::new(1024, 1024).with_bias(false).init(device);
let linear58 = LinearConfig::new(1024, 1024).with_bias(false).init(device);
let linear59 = LinearConfig::new(1024, 1024).with_bias(false).init(device);
let constant152: burn::module::Param<Tensor<1>> = burn::module::Param::uninitialized(
burn::module::ParamId::new(),
move |device, _require_grad| Tensor::<
1,
>::zeros([1024], (device, burn::tensor::DType::F32)),
device.clone(),
false,
[1024].into(),
);
let constant153: burn::module::Param<Tensor<1>> = burn::module::Param::uninitialized(
burn::module::ParamId::new(),
move |device, _require_grad| Tensor::<
1,
>::zeros([1024], (device, burn::tensor::DType::F32)),
device.clone(),
false,
[1024].into(),
);
let constant154: burn::module::Param<Tensor<1>> = burn::module::Param::uninitialized(
burn::module::ParamId::new(),
move |device, _require_grad| Tensor::<
1,
>::zeros([1024], (device, burn::tensor::DType::F32)),
device.clone(),
false,
[1024].into(),
);
let linear60 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
let layernormalization18 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let linear61 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
let linear62 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
let layernormalization19 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let linear63 = LinearConfig::new(1024, 4096).with_bias(true).init(device);
let linear64 = LinearConfig::new(4096, 1024).with_bias(true).init(device);
let layernormalization20 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let linear65 = LinearConfig::new(1024, 1024).with_bias(false).init(device);
let linear66 = LinearConfig::new(1024, 1024).with_bias(false).init(device);
let linear67 = LinearConfig::new(1024, 1024).with_bias(false).init(device);
let constant174: burn::module::Param<Tensor<1>> = burn::module::Param::uninitialized(
burn::module::ParamId::new(),
move |device, _require_grad| Tensor::<
1,
>::zeros([1024], (device, burn::tensor::DType::F32)),
device.clone(),
false,
[1024].into(),
);
let constant175: burn::module::Param<Tensor<1>> = burn::module::Param::uninitialized(
burn::module::ParamId::new(),
move |device, _require_grad| Tensor::<
1,
>::zeros([1024], (device, burn::tensor::DType::F32)),
device.clone(),
false,
[1024].into(),
);
let constant176: burn::module::Param<Tensor<1>> = burn::module::Param::uninitialized(
burn::module::ParamId::new(),
move |device, _require_grad| Tensor::<
1,
>::zeros([1024], (device, burn::tensor::DType::F32)),
device.clone(),
false,
[1024].into(),
);
let linear68 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
let layernormalization21 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let linear69 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
let linear70 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
let layernormalization22 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let linear71 = LinearConfig::new(1024, 4096).with_bias(true).init(device);
let linear72 = LinearConfig::new(4096, 1024).with_bias(true).init(device);
let layernormalization23 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let linear73 = LinearConfig::new(1024, 1024).with_bias(false).init(device);
let linear74 = LinearConfig::new(1024, 1024).with_bias(false).init(device);
let linear75 = LinearConfig::new(1024, 1024).with_bias(false).init(device);
let constant196: burn::module::Param<Tensor<1>> = burn::module::Param::uninitialized(
burn::module::ParamId::new(),
move |device, _require_grad| Tensor::<
1,
>::zeros([1024], (device, burn::tensor::DType::F32)),
device.clone(),
false,
[1024].into(),
);
let constant197: burn::module::Param<Tensor<1>> = burn::module::Param::uninitialized(
burn::module::ParamId::new(),
move |device, _require_grad| Tensor::<
1,
>::zeros([1024], (device, burn::tensor::DType::F32)),
device.clone(),
false,
[1024].into(),
);
let constant198: burn::module::Param<Tensor<1>> = burn::module::Param::uninitialized(
burn::module::ParamId::new(),
move |device, _require_grad| Tensor::<
1,
>::zeros([1024], (device, burn::tensor::DType::F32)),
device.clone(),
false,
[1024].into(),
);
let linear76 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
let layernormalization24 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let linear77 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
let linear78 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
let layernormalization25 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let linear79 = LinearConfig::new(1024, 4096).with_bias(true).init(device);
let linear80 = LinearConfig::new(4096, 1024).with_bias(true).init(device);
let layernormalization26 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let linear81 = LinearConfig::new(1024, 16384).with_bias(true).init(device);
Self {
layernormalization15,
linear53,
linear54,
layernormalization16,
linear55,
linear56,
layernormalization17,
linear57,
linear58,
linear59,
constant152,
constant153,
constant154,
linear60,
layernormalization18,
linear61,
linear62,
layernormalization19,
linear63,
linear64,
layernormalization20,
linear65,
linear66,
linear67,
constant174,
constant175,
constant176,
linear68,
layernormalization21,
linear69,
linear70,
layernormalization22,
linear71,
linear72,
layernormalization23,
linear73,
linear74,
linear75,
constant196,
constant197,
constant198,
linear76,
layernormalization24,
linear77,
linear78,
layernormalization25,
linear79,
linear80,
layernormalization26,
linear81,
device: device.clone(),
}
}
#[allow(clippy::let_and_return, clippy::approx_constant)]
pub fn forward(
&self,
add29_out1: Tensor<3>,
transpose9_out1: Tensor<4>,
transpose10_out1: Tensor<4>,
expand1_out1: Tensor<4>,
transpose11_out1: Tensor<4>,
transpose12_out1: Tensor<4>,
transpose13_out1: Tensor<4>,
transpose14_out1: Tensor<4>,
transpose15_out1: Tensor<4>,
transpose16_out1: Tensor<4>,
num_logits_to_keep: i64,
) -> (Tensor<3>, Tensor<4>, Tensor<4>, Tensor<4>, Tensor<4>, Tensor<4>, Tensor<4>) {
let layernormalization15_out1 = {
let dtype = add29_out1.clone().dtype();
self.layernormalization15
.forward(add29_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let linear53_out1 = self.linear53.forward(layernormalization15_out1);
let reshape45_out1 = linear53_out1.reshape([0, 0, 8, -1]);
let transpose45_out1 = reshape45_out1.permute([0, 2, 1, 3]);
let (attention10_out1,) = {
let q = transpose45_out1;
let k = transpose9_out1;
let v = transpose10_out1;
let attention10_out1 = burn::tensor::module::attention(
q,
k,
v,
None,
None,
burn::tensor::ops::AttentionModuleOptions {
scale: Some(0.0883883461356163f64),
softcap: None,
is_causal: false,
},
);
(attention10_out1,)
};
let transpose46_out1 = attention10_out1.permute([0, 2, 1, 3]);
let reshape46_out1 = transpose46_out1.reshape([0, 0, -1]);
let linear54_out1 = self.linear54.forward(reshape46_out1);
let add30_out1 = add29_out1.add(linear54_out1);
let layernormalization16_out1 = {
let dtype = add30_out1.clone().dtype();
self.layernormalization16
.forward(add30_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let linear55_out1 = self.linear55.forward(layernormalization16_out1);
let relu5_out1 = burn::tensor::activation::relu(linear55_out1);
let linear56_out1 = self.linear56.forward(relu5_out1);
let add31_out1 = add30_out1.add(linear56_out1);
let layernormalization17_out1 = {
let dtype = add31_out1.clone().dtype();
self.layernormalization17
.forward(add31_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let linear57_out1 = self.linear57.forward(layernormalization17_out1.clone());
let linear58_out1 = self.linear58.forward(layernormalization17_out1.clone());
let linear59_out1 = self.linear59.forward(layernormalization17_out1);
let constant152_out1 = self.constant152.val();
let add32_out1 = linear57_out1
.add((constant152_out1).unsqueeze_dims(&[0isize, 1isize]));
let constant153_out1 = self.constant153.val();
let add33_out1 = linear58_out1
.add((constant153_out1).unsqueeze_dims(&[0isize, 1isize]));
let constant154_out1 = self.constant154.val();
let add34_out1 = linear59_out1
.add((constant154_out1).unsqueeze_dims(&[0isize, 1isize]));
let reshape47_out1 = add33_out1.reshape([0, 0, 8, -1]);
let transpose47_out1 = reshape47_out1.permute([0, 2, 1, 3]);
let reshape48_out1 = add34_out1.reshape([0, 0, 8, -1]);
let transpose48_out1 = reshape48_out1.permute([0, 2, 1, 3]);
let reshape49_out1 = add32_out1.reshape([0, 0, 8, -1]);
let transpose49_out1 = reshape49_out1.permute([0, 2, 1, 3]);
let (attention11_out1,) = {
let q = transpose49_out1;
let k = transpose47_out1.clone();
let v = transpose48_out1.clone();
let hide = {
let [batch, heads, q_len, _] = q.dims();
let k_len = k.dims()[2];
let rows = Tensor::<
1,
Int,
>::arange(0..q_len as i64, (&self.device, burn::tensor::DType::I64))
.reshape([q_len, 1])
.expand([q_len, k_len]);
let cols = Tensor::<
1,
Int,
>::arange(0..k_len as i64, (&self.device, burn::tensor::DType::I64))
.reshape([1, k_len])
.expand([q_len, k_len]);
cols.greater(rows).unsqueeze::<4>().expand([batch, heads, q_len, k_len])
};
let score_bias = expand1_out1.clone();
let hidden_rows = {
let shape = [q.dims()[0], q.dims()[1], q.dims()[2], k.dims()[2]];
score_bias
.clone()
.expand(shape)
.mask_fill(hide.clone().expand(shape), f32::NEG_INFINITY)
.max_dim(3)
.equal_elem(f32::NEG_INFINITY)
};
let attention11_out1 = burn::tensor::module::attention(
q,
k,
v,
Some(hide),
Some(score_bias),
burn::tensor::ops::AttentionModuleOptions {
scale: Some(0.0883883461356163f64),
softcap: None,
is_causal: false,
},
);
let attention11_out1 = {
let dims = attention11_out1.dims();
attention11_out1.mask_fill(hidden_rows.expand(dims), 0.0)
};
(attention11_out1,)
};
let transpose50_out1 = attention11_out1.permute([0, 2, 1, 3]);
let reshape50_out1 = transpose50_out1.reshape([0, 0, -1]);
let linear60_out1 = self.linear60.forward(reshape50_out1);
let add35_out1 = add31_out1.add(linear60_out1);
let layernormalization18_out1 = {
let dtype = add35_out1.clone().dtype();
self.layernormalization18
.forward(add35_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let linear61_out1 = self.linear61.forward(layernormalization18_out1);
let reshape51_out1 = linear61_out1.reshape([0, 0, 8, -1]);
let transpose51_out1 = reshape51_out1.permute([0, 2, 1, 3]);
let (attention12_out1,) = {
let q = transpose51_out1;
let k = transpose11_out1;
let v = transpose12_out1;
let attention12_out1 = burn::tensor::module::attention(
q,
k,
v,
None,
None,
burn::tensor::ops::AttentionModuleOptions {
scale: Some(0.0883883461356163f64),
softcap: None,
is_causal: false,
},
);
(attention12_out1,)
};
let transpose52_out1 = attention12_out1.permute([0, 2, 1, 3]);
let reshape52_out1 = transpose52_out1.reshape([0, 0, -1]);
let linear62_out1 = self.linear62.forward(reshape52_out1);
let add36_out1 = add35_out1.add(linear62_out1);
let layernormalization19_out1 = {
let dtype = add36_out1.clone().dtype();
self.layernormalization19
.forward(add36_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let linear63_out1 = self.linear63.forward(layernormalization19_out1);
let relu6_out1 = burn::tensor::activation::relu(linear63_out1);
let linear64_out1 = self.linear64.forward(relu6_out1);
let add37_out1 = add36_out1.add(linear64_out1);
let layernormalization20_out1 = {
let dtype = add37_out1.clone().dtype();
self.layernormalization20
.forward(add37_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let linear65_out1 = self.linear65.forward(layernormalization20_out1.clone());
let linear66_out1 = self.linear66.forward(layernormalization20_out1.clone());
let linear67_out1 = self.linear67.forward(layernormalization20_out1);
let constant174_out1 = self.constant174.val();
let add38_out1 = linear65_out1
.add((constant174_out1).unsqueeze_dims(&[0isize, 1isize]));
let constant175_out1 = self.constant175.val();
let add39_out1 = linear66_out1
.add((constant175_out1).unsqueeze_dims(&[0isize, 1isize]));
let constant176_out1 = self.constant176.val();
let add40_out1 = linear67_out1
.add((constant176_out1).unsqueeze_dims(&[0isize, 1isize]));
let reshape53_out1 = add39_out1.reshape([0, 0, 8, -1]);
let transpose53_out1 = reshape53_out1.permute([0, 2, 1, 3]);
let reshape54_out1 = add40_out1.reshape([0, 0, 8, -1]);
let transpose54_out1 = reshape54_out1.permute([0, 2, 1, 3]);
let reshape55_out1 = add38_out1.reshape([0, 0, 8, -1]);
let transpose55_out1 = reshape55_out1.permute([0, 2, 1, 3]);
let (attention13_out1,) = {
let q = transpose55_out1;
let k = transpose53_out1.clone();
let v = transpose54_out1.clone();
let hide = {
let [batch, heads, q_len, _] = q.dims();
let k_len = k.dims()[2];
let rows = Tensor::<
1,
Int,
>::arange(0..q_len as i64, (&self.device, burn::tensor::DType::I64))
.reshape([q_len, 1])
.expand([q_len, k_len]);
let cols = Tensor::<
1,
Int,
>::arange(0..k_len as i64, (&self.device, burn::tensor::DType::I64))
.reshape([1, k_len])
.expand([q_len, k_len]);
cols.greater(rows).unsqueeze::<4>().expand([batch, heads, q_len, k_len])
};
let score_bias = expand1_out1.clone();
let hidden_rows = {
let shape = [q.dims()[0], q.dims()[1], q.dims()[2], k.dims()[2]];
score_bias
.clone()
.expand(shape)
.mask_fill(hide.clone().expand(shape), f32::NEG_INFINITY)
.max_dim(3)
.equal_elem(f32::NEG_INFINITY)
};
let attention13_out1 = burn::tensor::module::attention(
q,
k,
v,
Some(hide),
Some(score_bias),
burn::tensor::ops::AttentionModuleOptions {
scale: Some(0.0883883461356163f64),
softcap: None,
is_causal: false,
},
);
let attention13_out1 = {
let dims = attention13_out1.dims();
attention13_out1.mask_fill(hidden_rows.expand(dims), 0.0)
};
(attention13_out1,)
};
let transpose56_out1 = attention13_out1.permute([0, 2, 1, 3]);
let reshape56_out1 = transpose56_out1.reshape([0, 0, -1]);
let linear68_out1 = self.linear68.forward(reshape56_out1);
let add41_out1 = add37_out1.add(linear68_out1);
let layernormalization21_out1 = {
let dtype = add41_out1.clone().dtype();
self.layernormalization21
.forward(add41_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let linear69_out1 = self.linear69.forward(layernormalization21_out1);
let reshape57_out1 = linear69_out1.reshape([0, 0, 8, -1]);
let transpose57_out1 = reshape57_out1.permute([0, 2, 1, 3]);
let (attention14_out1,) = {
let q = transpose57_out1;
let k = transpose13_out1;
let v = transpose14_out1;
let attention14_out1 = burn::tensor::module::attention(
q,
k,
v,
None,
None,
burn::tensor::ops::AttentionModuleOptions {
scale: Some(0.0883883461356163f64),
softcap: None,
is_causal: false,
},
);
(attention14_out1,)
};
let transpose58_out1 = attention14_out1.permute([0, 2, 1, 3]);
let reshape58_out1 = transpose58_out1.reshape([0, 0, -1]);
let linear70_out1 = self.linear70.forward(reshape58_out1);
let add42_out1 = add41_out1.add(linear70_out1);
let layernormalization22_out1 = {
let dtype = add42_out1.clone().dtype();
self.layernormalization22
.forward(add42_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let linear71_out1 = self.linear71.forward(layernormalization22_out1);
let relu7_out1 = burn::tensor::activation::relu(linear71_out1);
let linear72_out1 = self.linear72.forward(relu7_out1);
let add43_out1 = add42_out1.add(linear72_out1);
let layernormalization23_out1 = {
let dtype = add43_out1.clone().dtype();
self.layernormalization23
.forward(add43_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let linear73_out1 = self.linear73.forward(layernormalization23_out1.clone());
let linear74_out1 = self.linear74.forward(layernormalization23_out1.clone());
let linear75_out1 = self.linear75.forward(layernormalization23_out1);
let constant196_out1 = self.constant196.val();
let add44_out1 = linear73_out1
.add((constant196_out1).unsqueeze_dims(&[0isize, 1isize]));
let constant197_out1 = self.constant197.val();
let add45_out1 = linear74_out1
.add((constant197_out1).unsqueeze_dims(&[0isize, 1isize]));
let constant198_out1 = self.constant198.val();
let add46_out1 = linear75_out1
.add((constant198_out1).unsqueeze_dims(&[0isize, 1isize]));
let reshape59_out1 = add45_out1.reshape([0, 0, 8, -1]);
let transpose59_out1 = reshape59_out1.permute([0, 2, 1, 3]);
let reshape60_out1 = add46_out1.reshape([0, 0, 8, -1]);
let transpose60_out1 = reshape60_out1.permute([0, 2, 1, 3]);
let reshape61_out1 = add44_out1.reshape([0, 0, 8, -1]);
let transpose61_out1 = reshape61_out1.permute([0, 2, 1, 3]);
let (attention15_out1,) = {
let q = transpose61_out1;
let k = transpose59_out1.clone();
let v = transpose60_out1.clone();
let hide = {
let [batch, heads, q_len, _] = q.dims();
let k_len = k.dims()[2];
let rows = Tensor::<
1,
Int,
>::arange(0..q_len as i64, (&self.device, burn::tensor::DType::I64))
.reshape([q_len, 1])
.expand([q_len, k_len]);
let cols = Tensor::<
1,
Int,
>::arange(0..k_len as i64, (&self.device, burn::tensor::DType::I64))
.reshape([1, k_len])
.expand([q_len, k_len]);
cols.greater(rows).unsqueeze::<4>().expand([batch, heads, q_len, k_len])
};
let score_bias = expand1_out1;
let hidden_rows = {
let shape = [q.dims()[0], q.dims()[1], q.dims()[2], k.dims()[2]];
score_bias
.clone()
.expand(shape)
.mask_fill(hide.clone().expand(shape), f32::NEG_INFINITY)
.max_dim(3)
.equal_elem(f32::NEG_INFINITY)
};
let attention15_out1 = burn::tensor::module::attention(
q,
k,
v,
Some(hide),
Some(score_bias),
burn::tensor::ops::AttentionModuleOptions {
scale: Some(0.0883883461356163f64),
softcap: None,
is_causal: false,
},
);
let attention15_out1 = {
let dims = attention15_out1.dims();
attention15_out1.mask_fill(hidden_rows.expand(dims), 0.0)
};
(attention15_out1,)
};
let transpose62_out1 = attention15_out1.permute([0, 2, 1, 3]);
let reshape62_out1 = transpose62_out1.reshape([0, 0, -1]);
let linear76_out1 = self.linear76.forward(reshape62_out1);
let add47_out1 = add43_out1.add(linear76_out1);
let layernormalization24_out1 = {
let dtype = add47_out1.clone().dtype();
self.layernormalization24
.forward(add47_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let linear77_out1 = self.linear77.forward(layernormalization24_out1);
let reshape63_out1 = linear77_out1.reshape([0, 0, 8, -1]);
let transpose63_out1 = reshape63_out1.permute([0, 2, 1, 3]);
let (attention16_out1,) = {
let q = transpose63_out1;
let k = transpose15_out1;
let v = transpose16_out1;
let attention16_out1 = burn::tensor::module::attention(
q,
k,
v,
None,
None,
burn::tensor::ops::AttentionModuleOptions {
scale: Some(0.0883883461356163f64),
softcap: None,
is_causal: false,
},
);
(attention16_out1,)
};
let transpose64_out1 = attention16_out1.permute([0, 2, 1, 3]);
let reshape64_out1 = transpose64_out1.reshape([0, 0, -1]);
let linear78_out1 = self.linear78.forward(reshape64_out1);
let add48_out1 = add47_out1.add(linear78_out1);
let layernormalization25_out1 = {
let dtype = add48_out1.clone().dtype();
self.layernormalization25
.forward(add48_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let linear79_out1 = self.linear79.forward(layernormalization25_out1);
let relu8_out1 = burn::tensor::activation::relu(linear79_out1);
let linear80_out1 = self.linear80.forward(relu8_out1);
let add49_out1 = add48_out1.add(linear80_out1);
let layernormalization26_out1 = {
let dtype = add49_out1.dtype();
self.layernormalization26
.forward(add49_out1.cast(burn::tensor::DType::F32))
.cast(dtype)
};
let neg1_out1 = -num_logits_to_keep;
let unsqueeze5_out1 = [neg1_out1 as i64];
let slice1_out1 = layernormalization26_out1
.slice(s![.., unsqueeze5_out1[0]..2147483647, ..]);
let linear81_out1 = self.linear81.forward(slice1_out1);
(
linear81_out1,
transpose47_out1,
transpose48_out1,
transpose53_out1,
transpose54_out1,
transpose59_out1,
transpose60_out1,
)
}
}
#[derive(Module, Debug)]
pub struct Model {
submodule1: Submodule1,
submodule2: Submodule2,
submodule3: Submodule3,
submodule4: Submodule4,
#[module(skip)]
device: Device,
}
impl Model {
#[allow(unused_variables)]
pub fn new(device: &Device) -> Self {
let submodule1 = Submodule1::new(device);
let submodule2 = Submodule2::new(device);
let submodule3 = Submodule3::new(device);
let submodule4 = Submodule4::new(device);
Self {
submodule1,
submodule2,
submodule3,
submodule4,
device: device.clone(),
}
}
#[allow(clippy::let_and_return, clippy::approx_constant)]
pub fn forward(
&self,
input_ids: Tensor<2, Int>,
attention_mask: Tensor<2, Int>,
position_ids: Tensor<2, Int>,
num_logits_to_keep: i64,
encoder_hidden_states: Tensor<3>,
) -> (
Tensor<3>,
Tensor<4>,
Tensor<4>,
Tensor<4>,
Tensor<4>,
Tensor<4>,
Tensor<4>,
Tensor<4>,
Tensor<4>,
Tensor<4>,
Tensor<4>,
Tensor<4>,
Tensor<4>,
Tensor<4>,
Tensor<4>,
Tensor<4>,
Tensor<4>,
Tensor<4>,
Tensor<4>,
Tensor<4>,
Tensor<4>,
Tensor<4>,
Tensor<4>,
Tensor<4>,
Tensor<4>,
Tensor<4>,
Tensor<4>,
Tensor<4>,
Tensor<4>,
Tensor<4>,
Tensor<4>,
Tensor<4>,
Tensor<4>,
) {
let (
add5_out1,
transpose1_out1,
transpose2_out1,
expand1_out1,
transpose3_out1,
transpose4_out1,
transpose5_out1,
transpose6_out1,
transpose7_out1,
transpose8_out1,
transpose9_out1,
transpose10_out1,
transpose11_out1,
transpose12_out1,
transpose13_out1,
transpose14_out1,
transpose15_out1,
transpose16_out1,
transpose17_out1,
transpose18_out1,
) = self
.submodule1
.forward(attention_mask, input_ids, position_ids, encoder_hidden_states);
let (
add17_out1,
transpose23_out1,
transpose24_out1,
transpose29_out1,
transpose30_out1,
) = self
.submodule2
.forward(
add5_out1,
transpose1_out1.clone(),
transpose2_out1.clone(),
expand1_out1.clone(),
transpose3_out1.clone(),
transpose4_out1.clone(),
);
let (
add29_out1,
transpose35_out1,
transpose36_out1,
transpose41_out1,
transpose42_out1,
) = self
.submodule3
.forward(
add17_out1,
transpose5_out1.clone(),
transpose6_out1.clone(),
expand1_out1.clone(),
transpose7_out1.clone(),
transpose8_out1.clone(),
);
let (
linear81_out1,
transpose47_out1,
transpose48_out1,
transpose53_out1,
transpose54_out1,
transpose59_out1,
transpose60_out1,
) = self
.submodule4
.forward(
add29_out1,
transpose9_out1.clone(),
transpose10_out1.clone(),
expand1_out1,
transpose11_out1.clone(),
transpose12_out1.clone(),
transpose13_out1.clone(),
transpose14_out1.clone(),
transpose15_out1.clone(),
transpose16_out1.clone(),
num_logits_to_keep,
);
(
linear81_out1,
transpose17_out1,
transpose18_out1,
transpose1_out1,
transpose2_out1,
transpose23_out1,
transpose24_out1,
transpose3_out1,
transpose4_out1,
transpose29_out1,
transpose30_out1,
transpose5_out1,
transpose6_out1,
transpose35_out1,
transpose36_out1,
transpose7_out1,
transpose8_out1,
transpose41_out1,
transpose42_out1,
transpose9_out1,
transpose10_out1,
transpose47_out1,
transpose48_out1,
transpose11_out1,
transpose12_out1,
transpose53_out1,
transpose54_out1,
transpose13_out1,
transpose14_out1,
transpose59_out1,
transpose60_out1,
transpose15_out1,
transpose16_out1,
)
}
}