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,
layernormalization2: LayerNorm,
linear1: Linear,
linear2: Linear,
linear3: Linear,
constant10: burn::module::Param<Tensor<1>>,
constant11: burn::module::Param<Tensor<1>>,
constant12: burn::module::Param<Tensor<1>>,
constant189: burn::module::Param<Tensor<1>>,
constant190: burn::module::Param<Tensor<1>>,
linear4: Linear,
layernormalization3: LayerNorm,
linear5: Linear,
linear6: Linear,
layernormalization4: LayerNorm,
linear7: Linear,
linear8: Linear,
layernormalization5: LayerNorm,
linear9: Linear,
linear10: Linear,
linear11: Linear,
constant32: burn::module::Param<Tensor<1>>,
constant33: burn::module::Param<Tensor<1>>,
constant34: burn::module::Param<Tensor<1>>,
linear12: Linear,
layernormalization6: LayerNorm,
linear13: Linear,
linear14: Linear,
layernormalization7: LayerNorm,
linear15: Linear,
linear16: Linear,
layernormalization8: LayerNorm,
linear17: Linear,
linear18: Linear,
linear19: Linear,
constant54: burn::module::Param<Tensor<1>>,
constant55: burn::module::Param<Tensor<1>>,
constant56: burn::module::Param<Tensor<1>>,
linear20: Linear,
layernormalization9: LayerNorm,
linear21: Linear,
linear22: Linear,
layernormalization10: LayerNorm,
linear23: Linear,
linear24: Linear,
layernormalization11: LayerNorm,
linear25: Linear,
linear26: Linear,
linear27: Linear,
constant76: burn::module::Param<Tensor<1>>,
constant77: burn::module::Param<Tensor<1>>,
constant78: burn::module::Param<Tensor<1>>,
linear28: Linear,
layernormalization12: LayerNorm,
linear29: Linear,
linear30: Linear,
layernormalization13: LayerNorm,
linear31: Linear,
linear32: Linear,
layernormalization14: LayerNorm,
linear33: Linear,
linear34: Linear,
linear35: Linear,
constant98: burn::module::Param<Tensor<1>>,
constant99: burn::module::Param<Tensor<1>>,
constant100: burn::module::Param<Tensor<1>>,
linear36: Linear,
layernormalization15: LayerNorm,
linear37: Linear,
linear38: Linear,
layernormalization16: LayerNorm,
linear39: Linear,
linear40: Linear,
layernormalization17: LayerNorm,
linear41: Linear,
linear42: Linear,
linear43: Linear,
constant120: burn::module::Param<Tensor<1>>,
constant121: burn::module::Param<Tensor<1>>,
constant122: burn::module::Param<Tensor<1>>,
linear44: Linear,
layernormalization18: LayerNorm,
linear45: Linear,
linear46: 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 layernormalization2 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let linear1 = LinearConfig::new(1024, 1024).with_bias(false).init(device);
let linear2 = LinearConfig::new(1024, 1024).with_bias(false).init(device);
let linear3 = LinearConfig::new(1024, 1024).with_bias(false).init(device);
let constant10: 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 constant11: 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 constant12: 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 constant189: 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 constant190: 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 linear4 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
let layernormalization3 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.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 layernormalization4 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let linear7 = LinearConfig::new(1024, 4096).with_bias(true).init(device);
let linear8 = LinearConfig::new(4096, 1024).with_bias(true).init(device);
let layernormalization5 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let linear9 = LinearConfig::new(1024, 1024).with_bias(false).init(device);
let linear10 = LinearConfig::new(1024, 1024).with_bias(false).init(device);
let linear11 = LinearConfig::new(1024, 1024).with_bias(false).init(device);
let constant32: 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 constant33: 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 constant34: 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 linear12 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
let layernormalization6 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.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 layernormalization7 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let linear15 = LinearConfig::new(1024, 4096).with_bias(true).init(device);
let linear16 = LinearConfig::new(4096, 1024).with_bias(true).init(device);
let layernormalization8 = 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 constant54: 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 constant55: 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 constant56: 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 linear20 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
let layernormalization9 = 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 layernormalization10 = 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 layernormalization11 = 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 constant76: 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 constant77: 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 constant78: 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 layernormalization12 = 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 layernormalization13 = 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 layernormalization14 = 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 constant98: 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 constant99: 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 constant100: 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);
let layernormalization15 = 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 layernormalization16 = 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 layernormalization17 = 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 constant120: 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 constant121: 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 constant122: 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 layernormalization18 = 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);
Self {
constant1,
constant2,
layernormalization1,
layernormalization2,
linear1,
linear2,
linear3,
constant10,
constant11,
constant12,
constant189,
constant190,
linear4,
layernormalization3,
linear5,
linear6,
layernormalization4,
linear7,
linear8,
layernormalization5,
linear9,
linear10,
linear11,
constant32,
constant33,
constant34,
linear12,
layernormalization6,
linear13,
linear14,
layernormalization7,
linear15,
linear16,
layernormalization8,
linear17,
linear18,
linear19,
constant54,
constant55,
constant56,
linear20,
layernormalization9,
linear21,
linear22,
layernormalization10,
linear23,
linear24,
layernormalization11,
linear25,
linear26,
linear27,
constant76,
constant77,
constant78,
linear28,
layernormalization12,
linear29,
linear30,
layernormalization13,
linear31,
linear32,
layernormalization14,
linear33,
linear34,
linear35,
constant98,
constant99,
constant100,
linear36,
layernormalization15,
linear37,
linear38,
layernormalization16,
linear39,
linear40,
layernormalization17,
linear41,
linear42,
linear43,
constant120,
constant121,
constant122,
linear44,
layernormalization18,
linear45,
linear46,
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>,
past_key_values_0_decoder_key: Tensor<4>,
past_key_values_0_decoder_value: Tensor<4>,
past_key_values_0_encoder_key: Tensor<4>,
past_key_values_0_encoder_value: Tensor<4>,
past_key_values_1_decoder_key: Tensor<4>,
past_key_values_1_decoder_value: Tensor<4>,
past_key_values_1_encoder_key: Tensor<4>,
past_key_values_1_encoder_value: Tensor<4>,
past_key_values_2_decoder_key: Tensor<4>,
past_key_values_2_decoder_value: Tensor<4>,
past_key_values_2_encoder_key: Tensor<4>,
past_key_values_2_encoder_value: Tensor<4>,
past_key_values_3_decoder_key: Tensor<4>,
past_key_values_3_decoder_value: Tensor<4>,
past_key_values_3_encoder_key: Tensor<4>,
past_key_values_3_encoder_value: Tensor<4>,
past_key_values_4_decoder_key: Tensor<4>,
past_key_values_4_decoder_value: Tensor<4>,
past_key_values_4_encoder_key: Tensor<4>,
past_key_values_4_encoder_value: Tensor<4>,
past_key_values_5_decoder_key: Tensor<4>,
past_key_values_5_decoder_value: Tensor<4>,
past_key_values_5_encoder_key: Tensor<4>,
past_key_values_5_encoder_value: Tensor<4>,
) -> (
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>,
) {
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 layernormalization2_out1 = {
let dtype = layernormalization1_out1.clone().dtype();
self.layernormalization2
.forward(layernormalization1_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let linear1_out1 = self.linear1.forward(layernormalization2_out1.clone());
let linear2_out1 = self.linear2.forward(layernormalization2_out1.clone());
let linear3_out1 = self.linear3.forward(layernormalization2_out1);
let constant10_out1 = self.constant10.val();
let add2_out1 = linear1_out1
.add((constant10_out1).unsqueeze_dims(&[0isize, 1isize]));
let constant11_out1 = self.constant11.val();
let add3_out1 = linear2_out1
.add((constant11_out1).unsqueeze_dims(&[0isize, 1isize]));
let constant12_out1 = self.constant12.val();
let add4_out1 = linear3_out1
.add((constant12_out1).unsqueeze_dims(&[0isize, 1isize]));
let cast1_out1 = attention_mask.float().cast(burn::tensor::DType::F32);
let constant189_out1 = self.constant189.val();
let sub1_out1 = (constant189_out1).unsqueeze_dims(&[0isize]).sub(cast1_out1);
let constant190_out1 = self.constant190.val();
let mul1_out1 = sub1_out1.mul((constant190_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 constant194_out1: [i64; 1] = [1i64];
let concat1_out1: [i64; 4usize] = [
&unsqueeze1_out1[..],
&constant194_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 (attention1_out1, attention1_out2, attention1_out3) = {
let q = add2_out1;
let k = add3_out1;
let v = add4_out1;
let [batch_size, q_sequence_length, q_hidden_size] = q.dims();
#[allow(clippy::identity_op)]
let head_size = q_hidden_size / 8usize;
let kv_sequence_length = k.dims()[1];
#[allow(clippy::identity_op)]
let v_head_size = v.dims()[2] / 8usize;
let q = q
.reshape([batch_size, q_sequence_length, 8usize, head_size])
.permute([0, 2, 1, 3]);
let k = k
.reshape([batch_size, kv_sequence_length, 8usize, head_size])
.permute([0, 2, 1, 3]);
let v = v
.reshape([batch_size, kv_sequence_length, 8usize, v_head_size])
.permute([0, 2, 1, 3]);
let attention1_out2 = Tensor::cat(
[past_key_values_0_decoder_key, k].to_vec(),
2,
);
let k = attention1_out2.clone();
let attention1_out3 = Tensor::cat(
[past_key_values_0_decoder_value, v].to_vec(),
2,
);
let v = attention1_out3.clone();
let score_bias = expand1_out1.clone();
let hidden_rows = score_bias
.clone()
.max_dim(3)
.equal_elem(f32::NEG_INFINITY);
let attention1_out1 = burn::tensor::module::attention(
q,
k,
v,
None,
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)
};
let attention1_out1 = attention1_out1
.permute([0, 2, 1, 3])
.reshape([batch_size as i32, q_sequence_length as i32, -1]);
(attention1_out1, attention1_out2, attention1_out3)
};
let linear4_out1 = self.linear4.forward(attention1_out1);
let add5_out1 = layernormalization1_out1.add(linear4_out1);
let layernormalization3_out1 = {
let dtype = add5_out1.clone().dtype();
self.layernormalization3
.forward(add5_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let linear5_out1 = self.linear5.forward(layernormalization3_out1);
let reshape1_out1 = linear5_out1.reshape([0, 0, 8, -1]);
let transpose1_out1 = reshape1_out1.permute([0, 2, 1, 3]);
let (attention2_out1,) = {
let q = transpose1_out1;
let k = past_key_values_0_encoder_key;
let v = past_key_values_0_encoder_value;
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 transpose2_out1 = attention2_out1.permute([0, 2, 1, 3]);
let reshape2_out1 = transpose2_out1.reshape([0, 0, -1]);
let linear6_out1 = self.linear6.forward(reshape2_out1);
let add6_out1 = add5_out1.add(linear6_out1);
let layernormalization4_out1 = {
let dtype = add6_out1.clone().dtype();
self.layernormalization4
.forward(add6_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let linear7_out1 = self.linear7.forward(layernormalization4_out1);
let relu1_out1 = burn::tensor::activation::relu(linear7_out1);
let linear8_out1 = self.linear8.forward(relu1_out1);
let add7_out1 = add6_out1.add(linear8_out1);
let layernormalization5_out1 = {
let dtype = add7_out1.clone().dtype();
self.layernormalization5
.forward(add7_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let linear9_out1 = self.linear9.forward(layernormalization5_out1.clone());
let linear10_out1 = self.linear10.forward(layernormalization5_out1.clone());
let linear11_out1 = self.linear11.forward(layernormalization5_out1);
let constant32_out1 = self.constant32.val();
let add8_out1 = linear9_out1
.add((constant32_out1).unsqueeze_dims(&[0isize, 1isize]));
let constant33_out1 = self.constant33.val();
let add9_out1 = linear10_out1
.add((constant33_out1).unsqueeze_dims(&[0isize, 1isize]));
let constant34_out1 = self.constant34.val();
let add10_out1 = linear11_out1
.add((constant34_out1).unsqueeze_dims(&[0isize, 1isize]));
let (attention3_out1, attention3_out2, attention3_out3) = {
let q = add8_out1;
let k = add9_out1;
let v = add10_out1;
let [batch_size, q_sequence_length, q_hidden_size] = q.dims();
#[allow(clippy::identity_op)]
let head_size = q_hidden_size / 8usize;
let kv_sequence_length = k.dims()[1];
#[allow(clippy::identity_op)]
let v_head_size = v.dims()[2] / 8usize;
let q = q
.reshape([batch_size, q_sequence_length, 8usize, head_size])
.permute([0, 2, 1, 3]);
let k = k
.reshape([batch_size, kv_sequence_length, 8usize, head_size])
.permute([0, 2, 1, 3]);
let v = v
.reshape([batch_size, kv_sequence_length, 8usize, v_head_size])
.permute([0, 2, 1, 3]);
let attention3_out2 = Tensor::cat(
[past_key_values_1_decoder_key, k].to_vec(),
2,
);
let k = attention3_out2.clone();
let attention3_out3 = Tensor::cat(
[past_key_values_1_decoder_value, v].to_vec(),
2,
);
let v = attention3_out3.clone();
let score_bias = expand1_out1.clone();
let hidden_rows = score_bias
.clone()
.max_dim(3)
.equal_elem(f32::NEG_INFINITY);
let attention3_out1 = burn::tensor::module::attention(
q,
k,
v,
None,
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)
};
let attention3_out1 = attention3_out1
.permute([0, 2, 1, 3])
.reshape([batch_size as i32, q_sequence_length as i32, -1]);
(attention3_out1, attention3_out2, attention3_out3)
};
let linear12_out1 = self.linear12.forward(attention3_out1);
let add11_out1 = add7_out1.add(linear12_out1);
let layernormalization6_out1 = {
let dtype = add11_out1.clone().dtype();
self.layernormalization6
.forward(add11_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let linear13_out1 = self.linear13.forward(layernormalization6_out1);
let reshape3_out1 = linear13_out1.reshape([0, 0, 8, -1]);
let transpose3_out1 = reshape3_out1.permute([0, 2, 1, 3]);
let (attention4_out1,) = {
let q = transpose3_out1;
let k = past_key_values_1_encoder_key;
let v = past_key_values_1_encoder_value;
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 transpose4_out1 = attention4_out1.permute([0, 2, 1, 3]);
let reshape4_out1 = transpose4_out1.reshape([0, 0, -1]);
let linear14_out1 = self.linear14.forward(reshape4_out1);
let add12_out1 = add11_out1.add(linear14_out1);
let layernormalization7_out1 = {
let dtype = add12_out1.clone().dtype();
self.layernormalization7
.forward(add12_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let linear15_out1 = self.linear15.forward(layernormalization7_out1);
let relu2_out1 = burn::tensor::activation::relu(linear15_out1);
let linear16_out1 = self.linear16.forward(relu2_out1);
let add13_out1 = add12_out1.add(linear16_out1);
let layernormalization8_out1 = {
let dtype = add13_out1.clone().dtype();
self.layernormalization8
.forward(add13_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let linear17_out1 = self.linear17.forward(layernormalization8_out1.clone());
let linear18_out1 = self.linear18.forward(layernormalization8_out1.clone());
let linear19_out1 = self.linear19.forward(layernormalization8_out1);
let constant54_out1 = self.constant54.val();
let add14_out1 = linear17_out1
.add((constant54_out1).unsqueeze_dims(&[0isize, 1isize]));
let constant55_out1 = self.constant55.val();
let add15_out1 = linear18_out1
.add((constant55_out1).unsqueeze_dims(&[0isize, 1isize]));
let constant56_out1 = self.constant56.val();
let add16_out1 = linear19_out1
.add((constant56_out1).unsqueeze_dims(&[0isize, 1isize]));
let (attention5_out1, attention5_out2, attention5_out3) = {
let q = add14_out1;
let k = add15_out1;
let v = add16_out1;
let [batch_size, q_sequence_length, q_hidden_size] = q.dims();
#[allow(clippy::identity_op)]
let head_size = q_hidden_size / 8usize;
let kv_sequence_length = k.dims()[1];
#[allow(clippy::identity_op)]
let v_head_size = v.dims()[2] / 8usize;
let q = q
.reshape([batch_size, q_sequence_length, 8usize, head_size])
.permute([0, 2, 1, 3]);
let k = k
.reshape([batch_size, kv_sequence_length, 8usize, head_size])
.permute([0, 2, 1, 3]);
let v = v
.reshape([batch_size, kv_sequence_length, 8usize, v_head_size])
.permute([0, 2, 1, 3]);
let attention5_out2 = Tensor::cat(
[past_key_values_2_decoder_key, k].to_vec(),
2,
);
let k = attention5_out2.clone();
let attention5_out3 = Tensor::cat(
[past_key_values_2_decoder_value, v].to_vec(),
2,
);
let v = attention5_out3.clone();
let score_bias = expand1_out1.clone();
let hidden_rows = score_bias
.clone()
.max_dim(3)
.equal_elem(f32::NEG_INFINITY);
let attention5_out1 = burn::tensor::module::attention(
q,
k,
v,
None,
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)
};
let attention5_out1 = attention5_out1
.permute([0, 2, 1, 3])
.reshape([batch_size as i32, q_sequence_length as i32, -1]);
(attention5_out1, attention5_out2, attention5_out3)
};
let linear20_out1 = self.linear20.forward(attention5_out1);
let add17_out1 = add13_out1.add(linear20_out1);
let layernormalization9_out1 = {
let dtype = add17_out1.clone().dtype();
self.layernormalization9
.forward(add17_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let linear21_out1 = self.linear21.forward(layernormalization9_out1);
let reshape5_out1 = linear21_out1.reshape([0, 0, 8, -1]);
let transpose5_out1 = reshape5_out1.permute([0, 2, 1, 3]);
let (attention6_out1,) = {
let q = transpose5_out1;
let k = past_key_values_2_encoder_key;
let v = past_key_values_2_encoder_value;
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 transpose6_out1 = attention6_out1.permute([0, 2, 1, 3]);
let reshape6_out1 = transpose6_out1.reshape([0, 0, -1]);
let linear22_out1 = self.linear22.forward(reshape6_out1);
let add18_out1 = add17_out1.add(linear22_out1);
let layernormalization10_out1 = {
let dtype = add18_out1.clone().dtype();
self.layernormalization10
.forward(add18_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let linear23_out1 = self.linear23.forward(layernormalization10_out1);
let relu3_out1 = burn::tensor::activation::relu(linear23_out1);
let linear24_out1 = self.linear24.forward(relu3_out1);
let add19_out1 = add18_out1.add(linear24_out1);
let layernormalization11_out1 = {
let dtype = add19_out1.clone().dtype();
self.layernormalization11
.forward(add19_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let linear25_out1 = self.linear25.forward(layernormalization11_out1.clone());
let linear26_out1 = self.linear26.forward(layernormalization11_out1.clone());
let linear27_out1 = self.linear27.forward(layernormalization11_out1);
let constant76_out1 = self.constant76.val();
let add20_out1 = linear25_out1
.add((constant76_out1).unsqueeze_dims(&[0isize, 1isize]));
let constant77_out1 = self.constant77.val();
let add21_out1 = linear26_out1
.add((constant77_out1).unsqueeze_dims(&[0isize, 1isize]));
let constant78_out1 = self.constant78.val();
let add22_out1 = linear27_out1
.add((constant78_out1).unsqueeze_dims(&[0isize, 1isize]));
let (attention7_out1, attention7_out2, attention7_out3) = {
let q = add20_out1;
let k = add21_out1;
let v = add22_out1;
let [batch_size, q_sequence_length, q_hidden_size] = q.dims();
#[allow(clippy::identity_op)]
let head_size = q_hidden_size / 8usize;
let kv_sequence_length = k.dims()[1];
#[allow(clippy::identity_op)]
let v_head_size = v.dims()[2] / 8usize;
let q = q
.reshape([batch_size, q_sequence_length, 8usize, head_size])
.permute([0, 2, 1, 3]);
let k = k
.reshape([batch_size, kv_sequence_length, 8usize, head_size])
.permute([0, 2, 1, 3]);
let v = v
.reshape([batch_size, kv_sequence_length, 8usize, v_head_size])
.permute([0, 2, 1, 3]);
let attention7_out2 = Tensor::cat(
[past_key_values_3_decoder_key, k].to_vec(),
2,
);
let k = attention7_out2.clone();
let attention7_out3 = Tensor::cat(
[past_key_values_3_decoder_value, v].to_vec(),
2,
);
let v = attention7_out3.clone();
let score_bias = expand1_out1.clone();
let hidden_rows = score_bias
.clone()
.max_dim(3)
.equal_elem(f32::NEG_INFINITY);
let attention7_out1 = burn::tensor::module::attention(
q,
k,
v,
None,
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)
};
let attention7_out1 = attention7_out1
.permute([0, 2, 1, 3])
.reshape([batch_size as i32, q_sequence_length as i32, -1]);
(attention7_out1, attention7_out2, attention7_out3)
};
let linear28_out1 = self.linear28.forward(attention7_out1);
let add23_out1 = add19_out1.add(linear28_out1);
let layernormalization12_out1 = {
let dtype = add23_out1.clone().dtype();
self.layernormalization12
.forward(add23_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let linear29_out1 = self.linear29.forward(layernormalization12_out1);
let reshape7_out1 = linear29_out1.reshape([0, 0, 8, -1]);
let transpose7_out1 = reshape7_out1.permute([0, 2, 1, 3]);
let (attention8_out1,) = {
let q = transpose7_out1;
let k = past_key_values_3_encoder_key;
let v = past_key_values_3_encoder_value;
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 transpose8_out1 = attention8_out1.permute([0, 2, 1, 3]);
let reshape8_out1 = transpose8_out1.reshape([0, 0, -1]);
let linear30_out1 = self.linear30.forward(reshape8_out1);
let add24_out1 = add23_out1.add(linear30_out1);
let layernormalization13_out1 = {
let dtype = add24_out1.clone().dtype();
self.layernormalization13
.forward(add24_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let linear31_out1 = self.linear31.forward(layernormalization13_out1);
let relu4_out1 = burn::tensor::activation::relu(linear31_out1);
let linear32_out1 = self.linear32.forward(relu4_out1);
let add25_out1 = add24_out1.add(linear32_out1);
let layernormalization14_out1 = {
let dtype = add25_out1.clone().dtype();
self.layernormalization14
.forward(add25_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let linear33_out1 = self.linear33.forward(layernormalization14_out1.clone());
let linear34_out1 = self.linear34.forward(layernormalization14_out1.clone());
let linear35_out1 = self.linear35.forward(layernormalization14_out1);
let constant98_out1 = self.constant98.val();
let add26_out1 = linear33_out1
.add((constant98_out1).unsqueeze_dims(&[0isize, 1isize]));
let constant99_out1 = self.constant99.val();
let add27_out1 = linear34_out1
.add((constant99_out1).unsqueeze_dims(&[0isize, 1isize]));
let constant100_out1 = self.constant100.val();
let add28_out1 = linear35_out1
.add((constant100_out1).unsqueeze_dims(&[0isize, 1isize]));
let (attention9_out1, attention9_out2, attention9_out3) = {
let q = add26_out1;
let k = add27_out1;
let v = add28_out1;
let [batch_size, q_sequence_length, q_hidden_size] = q.dims();
#[allow(clippy::identity_op)]
let head_size = q_hidden_size / 8usize;
let kv_sequence_length = k.dims()[1];
#[allow(clippy::identity_op)]
let v_head_size = v.dims()[2] / 8usize;
let q = q
.reshape([batch_size, q_sequence_length, 8usize, head_size])
.permute([0, 2, 1, 3]);
let k = k
.reshape([batch_size, kv_sequence_length, 8usize, head_size])
.permute([0, 2, 1, 3]);
let v = v
.reshape([batch_size, kv_sequence_length, 8usize, v_head_size])
.permute([0, 2, 1, 3]);
let attention9_out2 = Tensor::cat(
[past_key_values_4_decoder_key, k].to_vec(),
2,
);
let k = attention9_out2.clone();
let attention9_out3 = Tensor::cat(
[past_key_values_4_decoder_value, v].to_vec(),
2,
);
let v = attention9_out3.clone();
let score_bias = expand1_out1.clone();
let hidden_rows = score_bias
.clone()
.max_dim(3)
.equal_elem(f32::NEG_INFINITY);
let attention9_out1 = burn::tensor::module::attention(
q,
k,
v,
None,
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)
};
let attention9_out1 = attention9_out1
.permute([0, 2, 1, 3])
.reshape([batch_size as i32, q_sequence_length as i32, -1]);
(attention9_out1, attention9_out2, attention9_out3)
};
let linear36_out1 = self.linear36.forward(attention9_out1);
let add29_out1 = add25_out1.add(linear36_out1);
let layernormalization15_out1 = {
let dtype = add29_out1.clone().dtype();
self.layernormalization15
.forward(add29_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let linear37_out1 = self.linear37.forward(layernormalization15_out1);
let reshape9_out1 = linear37_out1.reshape([0, 0, 8, -1]);
let transpose9_out1 = reshape9_out1.permute([0, 2, 1, 3]);
let (attention10_out1,) = {
let q = transpose9_out1;
let k = past_key_values_4_encoder_key;
let v = past_key_values_4_encoder_value;
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 transpose10_out1 = attention10_out1.permute([0, 2, 1, 3]);
let reshape10_out1 = transpose10_out1.reshape([0, 0, -1]);
let linear38_out1 = self.linear38.forward(reshape10_out1);
let add30_out1 = add29_out1.add(linear38_out1);
let layernormalization16_out1 = {
let dtype = add30_out1.clone().dtype();
self.layernormalization16
.forward(add30_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let linear39_out1 = self.linear39.forward(layernormalization16_out1);
let relu5_out1 = burn::tensor::activation::relu(linear39_out1);
let linear40_out1 = self.linear40.forward(relu5_out1);
let add31_out1 = add30_out1.add(linear40_out1);
let layernormalization17_out1 = {
let dtype = add31_out1.clone().dtype();
self.layernormalization17
.forward(add31_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let linear41_out1 = self.linear41.forward(layernormalization17_out1.clone());
let linear42_out1 = self.linear42.forward(layernormalization17_out1.clone());
let linear43_out1 = self.linear43.forward(layernormalization17_out1);
let constant120_out1 = self.constant120.val();
let add32_out1 = linear41_out1
.add((constant120_out1).unsqueeze_dims(&[0isize, 1isize]));
let constant121_out1 = self.constant121.val();
let add33_out1 = linear42_out1
.add((constant121_out1).unsqueeze_dims(&[0isize, 1isize]));
let constant122_out1 = self.constant122.val();
let add34_out1 = linear43_out1
.add((constant122_out1).unsqueeze_dims(&[0isize, 1isize]));
let (attention11_out1, attention11_out2, attention11_out3) = {
let q = add32_out1;
let k = add33_out1;
let v = add34_out1;
let [batch_size, q_sequence_length, q_hidden_size] = q.dims();
#[allow(clippy::identity_op)]
let head_size = q_hidden_size / 8usize;
let kv_sequence_length = k.dims()[1];
#[allow(clippy::identity_op)]
let v_head_size = v.dims()[2] / 8usize;
let q = q
.reshape([batch_size, q_sequence_length, 8usize, head_size])
.permute([0, 2, 1, 3]);
let k = k
.reshape([batch_size, kv_sequence_length, 8usize, head_size])
.permute([0, 2, 1, 3]);
let v = v
.reshape([batch_size, kv_sequence_length, 8usize, v_head_size])
.permute([0, 2, 1, 3]);
let attention11_out2 = Tensor::cat(
[past_key_values_5_decoder_key, k].to_vec(),
2,
);
let k = attention11_out2.clone();
let attention11_out3 = Tensor::cat(
[past_key_values_5_decoder_value, v].to_vec(),
2,
);
let v = attention11_out3.clone();
let score_bias = expand1_out1.clone();
let hidden_rows = score_bias
.clone()
.max_dim(3)
.equal_elem(f32::NEG_INFINITY);
let attention11_out1 = burn::tensor::module::attention(
q,
k,
v,
None,
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)
};
let attention11_out1 = attention11_out1
.permute([0, 2, 1, 3])
.reshape([batch_size as i32, q_sequence_length as i32, -1]);
(attention11_out1, attention11_out2, attention11_out3)
};
let linear44_out1 = self.linear44.forward(attention11_out1);
let add35_out1 = add31_out1.add(linear44_out1);
let layernormalization18_out1 = {
let dtype = add35_out1.clone().dtype();
self.layernormalization18
.forward(add35_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let linear45_out1 = self.linear45.forward(layernormalization18_out1);
let reshape11_out1 = linear45_out1.reshape([0, 0, 8, -1]);
let transpose11_out1 = reshape11_out1.permute([0, 2, 1, 3]);
let (attention12_out1,) = {
let q = transpose11_out1;
let k = past_key_values_5_encoder_key;
let v = past_key_values_5_encoder_value;
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 transpose12_out1 = attention12_out1.permute([0, 2, 1, 3]);
let reshape12_out1 = transpose12_out1.reshape([0, 0, -1]);
let linear46_out1 = self.linear46.forward(reshape12_out1);
let add36_out1 = add35_out1.add(linear46_out1);
(
add36_out1,
expand1_out1,
attention1_out2,
attention1_out3,
attention3_out2,
attention3_out3,
attention5_out2,
attention5_out3,
attention7_out2,
attention7_out3,
attention9_out2,
attention9_out3,
attention11_out2,
attention11_out3,
)
}
}
#[derive(Module, Debug)]
pub struct Submodule2 {
layernormalization19: LayerNorm,
linear47: Linear,
linear48: Linear,
layernormalization20: LayerNorm,
linear49: Linear,
linear50: Linear,
linear51: Linear,
constant142: burn::module::Param<Tensor<1>>,
constant143: burn::module::Param<Tensor<1>>,
constant144: burn::module::Param<Tensor<1>>,
linear52: Linear,
layernormalization21: LayerNorm,
linear53: Linear,
linear54: Linear,
layernormalization22: LayerNorm,
linear55: Linear,
linear56: Linear,
layernormalization23: LayerNorm,
linear57: Linear,
linear58: Linear,
linear59: Linear,
constant164: burn::module::Param<Tensor<1>>,
constant165: burn::module::Param<Tensor<1>>,
constant166: burn::module::Param<Tensor<1>>,
linear60: Linear,
layernormalization24: LayerNorm,
linear61: Linear,
linear62: Linear,
layernormalization25: LayerNorm,
linear63: Linear,
linear64: Linear,
layernormalization26: LayerNorm,
linear65: Linear,
#[module(skip)]
device: Device,
}
impl Submodule2 {
#[allow(unused_variables)]
pub fn new(device: &Device) -> Self {
let layernormalization19 = 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 layernormalization20 = 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 constant142: 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 constant143: 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 constant144: 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);
let layernormalization21 = 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 layernormalization22 = 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 layernormalization23 = 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 constant164: 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 constant165: 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 constant166: 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 layernormalization24 = 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 layernormalization25 = 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 layernormalization26 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let linear65 = LinearConfig::new(1024, 16384).with_bias(true).init(device);
Self {
layernormalization19,
linear47,
linear48,
layernormalization20,
linear49,
linear50,
linear51,
constant142,
constant143,
constant144,
linear52,
layernormalization21,
linear53,
linear54,
layernormalization22,
linear55,
linear56,
layernormalization23,
linear57,
linear58,
linear59,
constant164,
constant165,
constant166,
linear60,
layernormalization24,
linear61,
linear62,
layernormalization25,
linear63,
linear64,
layernormalization26,
linear65,
device: device.clone(),
}
}
#[allow(clippy::let_and_return, clippy::approx_constant)]
pub fn forward(
&self,
add36_out1: Tensor<3>,
expand1_out1: Tensor<4>,
past_key_values_6_decoder_key: Tensor<4>,
past_key_values_6_decoder_value: Tensor<4>,
past_key_values_6_encoder_key: Tensor<4>,
past_key_values_6_encoder_value: Tensor<4>,
past_key_values_7_decoder_key: Tensor<4>,
past_key_values_7_decoder_value: Tensor<4>,
past_key_values_7_encoder_key: Tensor<4>,
past_key_values_7_encoder_value: Tensor<4>,
num_logits_to_keep: i64,
) -> (Tensor<3>, Tensor<4>, Tensor<4>, Tensor<4>, Tensor<4>) {
let layernormalization19_out1 = {
let dtype = add36_out1.clone().dtype();
self.layernormalization19
.forward(add36_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let linear47_out1 = self.linear47.forward(layernormalization19_out1);
let relu6_out1 = burn::tensor::activation::relu(linear47_out1);
let linear48_out1 = self.linear48.forward(relu6_out1);
let add37_out1 = add36_out1.add(linear48_out1);
let layernormalization20_out1 = {
let dtype = add37_out1.clone().dtype();
self.layernormalization20
.forward(add37_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let linear49_out1 = self.linear49.forward(layernormalization20_out1.clone());
let linear50_out1 = self.linear50.forward(layernormalization20_out1.clone());
let linear51_out1 = self.linear51.forward(layernormalization20_out1);
let constant142_out1 = self.constant142.val();
let add38_out1 = linear49_out1
.add((constant142_out1).unsqueeze_dims(&[0isize, 1isize]));
let constant143_out1 = self.constant143.val();
let add39_out1 = linear50_out1
.add((constant143_out1).unsqueeze_dims(&[0isize, 1isize]));
let constant144_out1 = self.constant144.val();
let add40_out1 = linear51_out1
.add((constant144_out1).unsqueeze_dims(&[0isize, 1isize]));
let (attention13_out1, attention13_out2, attention13_out3) = {
let q = add38_out1;
let k = add39_out1;
let v = add40_out1;
let [batch_size, q_sequence_length, q_hidden_size] = q.dims();
#[allow(clippy::identity_op)]
let head_size = q_hidden_size / 8usize;
let kv_sequence_length = k.dims()[1];
#[allow(clippy::identity_op)]
let v_head_size = v.dims()[2] / 8usize;
let q = q
.reshape([batch_size, q_sequence_length, 8usize, head_size])
.permute([0, 2, 1, 3]);
let k = k
.reshape([batch_size, kv_sequence_length, 8usize, head_size])
.permute([0, 2, 1, 3]);
let v = v
.reshape([batch_size, kv_sequence_length, 8usize, v_head_size])
.permute([0, 2, 1, 3]);
let attention13_out2 = Tensor::cat(
[past_key_values_6_decoder_key, k].to_vec(),
2,
);
let k = attention13_out2.clone();
let attention13_out3 = Tensor::cat(
[past_key_values_6_decoder_value, v].to_vec(),
2,
);
let v = attention13_out3.clone();
let score_bias = expand1_out1.clone();
let hidden_rows = score_bias
.clone()
.max_dim(3)
.equal_elem(f32::NEG_INFINITY);
let attention13_out1 = burn::tensor::module::attention(
q,
k,
v,
None,
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)
};
let attention13_out1 = attention13_out1
.permute([0, 2, 1, 3])
.reshape([batch_size as i32, q_sequence_length as i32, -1]);
(attention13_out1, attention13_out2, attention13_out3)
};
let linear52_out1 = self.linear52.forward(attention13_out1);
let add41_out1 = add37_out1.add(linear52_out1);
let layernormalization21_out1 = {
let dtype = add41_out1.clone().dtype();
self.layernormalization21
.forward(add41_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let linear53_out1 = self.linear53.forward(layernormalization21_out1);
let reshape13_out1 = linear53_out1.reshape([0, 0, 8, -1]);
let transpose13_out1 = reshape13_out1.permute([0, 2, 1, 3]);
let (attention14_out1,) = {
let q = transpose13_out1;
let k = past_key_values_6_encoder_key;
let v = past_key_values_6_encoder_value;
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 transpose14_out1 = attention14_out1.permute([0, 2, 1, 3]);
let reshape14_out1 = transpose14_out1.reshape([0, 0, -1]);
let linear54_out1 = self.linear54.forward(reshape14_out1);
let add42_out1 = add41_out1.add(linear54_out1);
let layernormalization22_out1 = {
let dtype = add42_out1.clone().dtype();
self.layernormalization22
.forward(add42_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let linear55_out1 = self.linear55.forward(layernormalization22_out1);
let relu7_out1 = burn::tensor::activation::relu(linear55_out1);
let linear56_out1 = self.linear56.forward(relu7_out1);
let add43_out1 = add42_out1.add(linear56_out1);
let layernormalization23_out1 = {
let dtype = add43_out1.clone().dtype();
self.layernormalization23
.forward(add43_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let linear57_out1 = self.linear57.forward(layernormalization23_out1.clone());
let linear58_out1 = self.linear58.forward(layernormalization23_out1.clone());
let linear59_out1 = self.linear59.forward(layernormalization23_out1);
let constant164_out1 = self.constant164.val();
let add44_out1 = linear57_out1
.add((constant164_out1).unsqueeze_dims(&[0isize, 1isize]));
let constant165_out1 = self.constant165.val();
let add45_out1 = linear58_out1
.add((constant165_out1).unsqueeze_dims(&[0isize, 1isize]));
let constant166_out1 = self.constant166.val();
let add46_out1 = linear59_out1
.add((constant166_out1).unsqueeze_dims(&[0isize, 1isize]));
let (attention15_out1, attention15_out2, attention15_out3) = {
let q = add44_out1;
let k = add45_out1;
let v = add46_out1;
let [batch_size, q_sequence_length, q_hidden_size] = q.dims();
#[allow(clippy::identity_op)]
let head_size = q_hidden_size / 8usize;
let kv_sequence_length = k.dims()[1];
#[allow(clippy::identity_op)]
let v_head_size = v.dims()[2] / 8usize;
let q = q
.reshape([batch_size, q_sequence_length, 8usize, head_size])
.permute([0, 2, 1, 3]);
let k = k
.reshape([batch_size, kv_sequence_length, 8usize, head_size])
.permute([0, 2, 1, 3]);
let v = v
.reshape([batch_size, kv_sequence_length, 8usize, v_head_size])
.permute([0, 2, 1, 3]);
let attention15_out2 = Tensor::cat(
[past_key_values_7_decoder_key, k].to_vec(),
2,
);
let k = attention15_out2.clone();
let attention15_out3 = Tensor::cat(
[past_key_values_7_decoder_value, v].to_vec(),
2,
);
let v = attention15_out3.clone();
let score_bias = expand1_out1;
let hidden_rows = score_bias
.clone()
.max_dim(3)
.equal_elem(f32::NEG_INFINITY);
let attention15_out1 = burn::tensor::module::attention(
q,
k,
v,
None,
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)
};
let attention15_out1 = attention15_out1
.permute([0, 2, 1, 3])
.reshape([batch_size as i32, q_sequence_length as i32, -1]);
(attention15_out1, attention15_out2, attention15_out3)
};
let linear60_out1 = self.linear60.forward(attention15_out1);
let add47_out1 = add43_out1.add(linear60_out1);
let layernormalization24_out1 = {
let dtype = add47_out1.clone().dtype();
self.layernormalization24
.forward(add47_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let linear61_out1 = self.linear61.forward(layernormalization24_out1);
let reshape15_out1 = linear61_out1.reshape([0, 0, 8, -1]);
let transpose15_out1 = reshape15_out1.permute([0, 2, 1, 3]);
let (attention16_out1,) = {
let q = transpose15_out1;
let k = past_key_values_7_encoder_key;
let v = past_key_values_7_encoder_value;
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 transpose16_out1 = attention16_out1.permute([0, 2, 1, 3]);
let reshape16_out1 = transpose16_out1.reshape([0, 0, -1]);
let linear62_out1 = self.linear62.forward(reshape16_out1);
let add48_out1 = add47_out1.add(linear62_out1);
let layernormalization25_out1 = {
let dtype = add48_out1.clone().dtype();
self.layernormalization25
.forward(add48_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let linear63_out1 = self.linear63.forward(layernormalization25_out1);
let relu8_out1 = burn::tensor::activation::relu(linear63_out1);
let linear64_out1 = self.linear64.forward(relu8_out1);
let add49_out1 = add48_out1.add(linear64_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 linear65_out1 = self.linear65.forward(slice1_out1);
(
linear65_out1,
attention13_out2,
attention13_out3,
attention15_out2,
attention15_out3,
)
}
}
#[derive(Module, Debug)]
pub struct Model {
submodule1: Submodule1,
submodule2: Submodule2,
#[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);
Self {
submodule1,
submodule2,
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,
past_key_values_0_decoder_key: Tensor<4>,
past_key_values_0_decoder_value: Tensor<4>,
past_key_values_0_encoder_key: Tensor<4>,
past_key_values_0_encoder_value: Tensor<4>,
past_key_values_1_decoder_key: Tensor<4>,
past_key_values_1_decoder_value: Tensor<4>,
past_key_values_1_encoder_key: Tensor<4>,
past_key_values_1_encoder_value: Tensor<4>,
past_key_values_2_decoder_key: Tensor<4>,
past_key_values_2_decoder_value: Tensor<4>,
past_key_values_2_encoder_key: Tensor<4>,
past_key_values_2_encoder_value: Tensor<4>,
past_key_values_3_decoder_key: Tensor<4>,
past_key_values_3_decoder_value: Tensor<4>,
past_key_values_3_encoder_key: Tensor<4>,
past_key_values_3_encoder_value: Tensor<4>,
past_key_values_4_decoder_key: Tensor<4>,
past_key_values_4_decoder_value: Tensor<4>,
past_key_values_4_encoder_key: Tensor<4>,
past_key_values_4_encoder_value: Tensor<4>,
past_key_values_5_decoder_key: Tensor<4>,
past_key_values_5_decoder_value: Tensor<4>,
past_key_values_5_encoder_key: Tensor<4>,
past_key_values_5_encoder_value: Tensor<4>,
past_key_values_6_decoder_key: Tensor<4>,
past_key_values_6_decoder_value: Tensor<4>,
past_key_values_6_encoder_key: Tensor<4>,
past_key_values_6_encoder_value: Tensor<4>,
past_key_values_7_decoder_key: Tensor<4>,
past_key_values_7_decoder_value: Tensor<4>,
past_key_values_7_encoder_key: Tensor<4>,
past_key_values_7_encoder_value: Tensor<4>,
) -> (
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>,
) {
let (
add36_out1,
expand1_out1,
attention1_out2,
attention1_out3,
attention3_out2,
attention3_out3,
attention5_out2,
attention5_out3,
attention7_out2,
attention7_out3,
attention9_out2,
attention9_out3,
attention11_out2,
attention11_out3,
) = self
.submodule1
.forward(
attention_mask,
input_ids,
position_ids,
past_key_values_0_decoder_key,
past_key_values_0_decoder_value,
past_key_values_0_encoder_key,
past_key_values_0_encoder_value,
past_key_values_1_decoder_key,
past_key_values_1_decoder_value,
past_key_values_1_encoder_key,
past_key_values_1_encoder_value,
past_key_values_2_decoder_key,
past_key_values_2_decoder_value,
past_key_values_2_encoder_key,
past_key_values_2_encoder_value,
past_key_values_3_decoder_key,
past_key_values_3_decoder_value,
past_key_values_3_encoder_key,
past_key_values_3_encoder_value,
past_key_values_4_decoder_key,
past_key_values_4_decoder_value,
past_key_values_4_encoder_key,
past_key_values_4_encoder_value,
past_key_values_5_decoder_key,
past_key_values_5_decoder_value,
past_key_values_5_encoder_key,
past_key_values_5_encoder_value,
);
let (
linear65_out1,
attention13_out2,
attention13_out3,
attention15_out2,
attention15_out3,
) = self
.submodule2
.forward(
add36_out1,
expand1_out1,
past_key_values_6_decoder_key,
past_key_values_6_decoder_value,
past_key_values_6_encoder_key,
past_key_values_6_encoder_value,
past_key_values_7_decoder_key,
past_key_values_7_decoder_value,
past_key_values_7_encoder_key,
past_key_values_7_encoder_value,
num_logits_to_keep,
);
(
linear65_out1,
attention1_out2,
attention1_out3,
attention3_out2,
attention3_out3,
attention5_out2,
attention5_out3,
attention7_out2,
attention7_out3,
attention9_out2,
attention9_out3,
attention11_out2,
attention11_out3,
attention13_out2,
attention13_out3,
attention15_out2,
attention15_out3,
)
}
}