extern crate alloc;
use burn::prelude::*;
use burn::nn::LayerNorm;
use burn::nn::LayerNormConfig;
use burn::nn::Linear;
use burn::nn::LinearConfig;
use burn::nn::PaddingConfig1d;
use burn::nn::PaddingConfig2d;
use burn::nn::conv::Conv1d;
use burn::nn::conv::Conv1dConfig;
use burn::nn::conv::Conv2d;
use burn::nn::conv::Conv2dConfig;
#[derive(Module, Debug)]
pub struct Submodule1 {
conv2d1: Conv2d,
constant930: burn::module::Param<Tensor<1>>,
constant931: burn::module::Param<Tensor<1>>,
constant932: burn::module::Param<Tensor<1>>,
constant933: burn::module::Param<Tensor<1>>,
constant934: burn::module::Param<Tensor<1>>,
#[module(skip)]
device: Device,
}
impl Submodule1 {
#[allow(unused_variables)]
pub fn new(device: &Device) -> Self {
let conv2d1 = Conv2dConfig::new([1, 256], [3, 3])
.with_stride([2, 2])
.with_padding(PaddingConfig2d::Explicit(1, 1, 1, 1))
.with_dilation([1, 1])
.with_groups(1)
.with_bias(true)
.init(device);
let constant930: 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 constant931: 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 constant932: 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([3f64]),
(device, burn::tensor::DType::F32),
),
device.clone(),
false,
[1].into(),
);
let constant933: 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([2f64]),
(device, burn::tensor::DType::F32),
),
device.clone(),
false,
[1].into(),
);
let constant934: 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(),
);
Self {
conv2d1,
constant930,
constant931,
constant932,
constant933,
constant934,
device: device.clone(),
}
}
#[allow(clippy::let_and_return, clippy::approx_constant)]
pub fn forward(
&self,
audio_signal: Tensor<3>,
length: Tensor<1, Int>,
) -> (Tensor<4>, Tensor<1>) {
let transpose1_out1 = audio_signal.permute([0, 2, 1]);
let unsqueeze1_out1: Tensor<4> = transpose1_out1.unsqueeze_dims::<4>(&[1]);
let cast1_out1 = length.float().cast(burn::tensor::DType::F32);
let cast2_out1 = cast1_out1.clone().int().cast(burn::tensor::DType::I64);
let shape1_out1: [i64; 4] = {
let axes = &unsqueeze1_out1.clone().dims()[0..4];
let mut output = [0i64; 4];
for i in 0..4 {
output[i] = axes[i] as i64;
}
output
};
let gather1_out1 = shape1_out1[0] as i64;
let gather2_out1 = shape1_out1[2] as i64;
let gather3_out1 = shape1_out1[3] as i64;
let range1_out1 = {
let start = 0i64;
let delta = 1i64;
assert!(delta != 0);
let n = (((gather2_out1 as i128 - start as i128) as f64) / delta as f64)
.ceil()
.max(0.0) as i64;
Tensor::arange(0..n, &self.device)
.cast(burn::tensor::DType::I64)
.mul_scalar(delta)
.add_scalar(start)
};
let unsqueeze2_out1 = [gather1_out1 as i64];
let unsqueeze3_out1 = [gather2_out1 as i64];
let concat1_out1: [i64; 2usize] = [&unsqueeze2_out1[..], &unsqueeze3_out1[..]]
.concat()
.try_into()
.unwrap();
let reshape1_out1 = concat1_out1;
let constant912_out1 = -1i64;
let mul2_out1 = {
let mut result = reshape1_out1;
let scalar = constant912_out1;
for result_item in result.iter_mut() {
*result_item = result_item.saturating_mul(scalar);
}
result
};
let max1_out1 = {
let __lhs = reshape1_out1;
let __rhs = mul2_out1;
core::array::from_fn::<i64, 2usize, _>(|__i| (__lhs[__i]).max(__rhs[__i]))
};
let expand1_out1 = {
let onnx_shape: [i64; 2usize] = max1_out1;
let input_dims = range1_out1.dims();
let mut shape = onnx_shape;
#[allow(clippy::needless_range_loop)]
for i in 0..1usize {
let dim_offset = 1usize + i;
if shape[dim_offset] == 1 && input_dims[i] > 1 {
shape[dim_offset] = input_dims[i] as i64;
}
}
range1_out1.expand(shape)
};
let unsqueeze4_out1: Tensor<2, Int> = cast2_out1.unsqueeze_dims::<2>(&[1]);
let less1_out1 = expand1_out1.lower(unsqueeze4_out1);
let unsqueeze5_out1: Tensor<3, Bool> = less1_out1.unsqueeze_dims::<3>(&[-1]);
let unsqueeze6_out1 = [gather1_out1 as i64];
let unsqueeze7_out1 = [gather2_out1 as i64];
let unsqueeze8_out1 = [gather3_out1 as i64];
let concat2_out1: [i64; 3usize] = [
&unsqueeze6_out1[..],
&unsqueeze7_out1[..],
&unsqueeze8_out1[..],
]
.concat()
.try_into()
.unwrap();
let reshape2_out1 = concat2_out1;
let constant920_out1 = -1i64;
let mul4_out1 = {
let mut result = reshape2_out1;
let scalar = constant920_out1;
for result_item in result.iter_mut() {
*result_item = result_item.saturating_mul(scalar);
}
result
};
let max2_out1 = {
let __lhs = reshape2_out1;
let __rhs = mul4_out1;
core::array::from_fn::<i64, 3usize, _>(|__i| (__lhs[__i]).max(__rhs[__i]))
};
let expand2_out1 = {
let onnx_shape: [i64; 3usize] = max2_out1;
let input_dims = unsqueeze5_out1.dims();
let mut shape = onnx_shape;
#[allow(clippy::needless_range_loop)]
for i in 0..3usize {
let dim_offset = i;
if shape[dim_offset] == 1 && input_dims[i] > 1 {
shape[dim_offset] = input_dims[i] as i64;
}
}
unsqueeze5_out1.expand(shape)
};
let cast4_out1 = expand2_out1.float().cast(burn::tensor::DType::F32);
let gather4_out1 = shape1_out1[1] as i64;
let unsqueeze9_out1: Tensor<4> = cast4_out1.unsqueeze_dims::<4>(&[1]);
let unsqueeze10_out1 = [gather1_out1 as i64];
let unsqueeze11_out1 = [gather4_out1 as i64];
let unsqueeze12_out1 = [gather2_out1 as i64];
let unsqueeze13_out1 = [gather3_out1 as i64];
let concat3_out1: [i64; 4usize] = [
&unsqueeze10_out1[..],
&unsqueeze11_out1[..],
&unsqueeze12_out1[..],
&unsqueeze13_out1[..],
]
.concat()
.try_into()
.unwrap();
let reshape3_out1 = concat3_out1;
let constant929_out1 = -1i64;
let mul6_out1 = {
let mut result = reshape3_out1;
let scalar = constant929_out1;
for result_item in result.iter_mut() {
*result_item = result_item.saturating_mul(scalar);
}
result
};
let max3_out1 = {
let __lhs = reshape3_out1;
let __rhs = mul6_out1;
core::array::from_fn::<i64, 4usize, _>(|__i| (__lhs[__i]).max(__rhs[__i]))
};
let expand3_out1 = {
let onnx_shape: [i64; 4usize] = max3_out1;
let input_dims = unsqueeze9_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;
}
}
unsqueeze9_out1.expand(shape)
};
let mul7_out1 = unsqueeze1_out1.mul(expand3_out1);
let conv2d1_out1 = self.conv2d1.forward(mul7_out1);
let constant930_out1 = self.constant930.val();
let add1_out1 = cast1_out1.add(constant930_out1);
let constant931_out1 = self.constant931.val();
let add2_out1 = add1_out1.add(constant931_out1);
let constant932_out1 = self.constant932.val();
let sub1_out1 = add2_out1.sub(constant932_out1);
let constant933_out1 = self.constant933.val();
let div1_out1 = sub1_out1.div(constant933_out1);
let cast5_out1 = div1_out1.int().cast(burn::tensor::DType::I64);
let cast6_out1 = cast5_out1.float().cast(burn::tensor::DType::F32);
let constant934_out1 = self.constant934.val();
let add3_out1 = cast6_out1.add(constant934_out1);
let cast7_out1 = add3_out1.clone().int().cast(burn::tensor::DType::I64);
let shape8_out1: [i64; 4] = {
let axes = &conv2d1_out1.clone().dims()[0..4];
let mut output = [0i64; 4];
for i in 0..4 {
output[i] = axes[i] as i64;
}
output
};
let gather5_out1 = shape8_out1[0] as i64;
let gather6_out1 = shape8_out1[2] as i64;
let gather7_out1 = shape8_out1[3] as i64;
let range2_out1 = {
let start = 0i64;
let delta = 1i64;
assert!(delta != 0);
let n = (((gather6_out1 as i128 - start as i128) as f64) / delta as f64)
.ceil()
.max(0.0) as i64;
Tensor::arange(0..n, &self.device)
.cast(burn::tensor::DType::I64)
.mul_scalar(delta)
.add_scalar(start)
};
let unsqueeze14_out1 = [gather5_out1 as i64];
let unsqueeze15_out1 = [gather6_out1 as i64];
let concat4_out1: [i64; 2usize] = [&unsqueeze14_out1[..], &unsqueeze15_out1[..]]
.concat()
.try_into()
.unwrap();
let reshape4_out1 = concat4_out1;
let constant944_out1 = -1i64;
let mul9_out1 = {
let mut result = reshape4_out1;
let scalar = constant944_out1;
for result_item in result.iter_mut() {
*result_item = result_item.saturating_mul(scalar);
}
result
};
let max4_out1 = {
let __lhs = reshape4_out1;
let __rhs = mul9_out1;
core::array::from_fn::<i64, 2usize, _>(|__i| (__lhs[__i]).max(__rhs[__i]))
};
let expand4_out1 = {
let onnx_shape: [i64; 2usize] = max4_out1;
let input_dims = range2_out1.dims();
let mut shape = onnx_shape;
#[allow(clippy::needless_range_loop)]
for i in 0..1usize {
let dim_offset = 1usize + i;
if shape[dim_offset] == 1 && input_dims[i] > 1 {
shape[dim_offset] = input_dims[i] as i64;
}
}
range2_out1.expand(shape)
};
let unsqueeze16_out1: Tensor<2, Int> = cast7_out1.unsqueeze_dims::<2>(&[1]);
let less2_out1 = expand4_out1.lower(unsqueeze16_out1);
let unsqueeze17_out1: Tensor<3, Bool> = less2_out1.unsqueeze_dims::<3>(&[-1]);
let unsqueeze18_out1 = [gather5_out1 as i64];
let unsqueeze19_out1 = [gather6_out1 as i64];
let unsqueeze20_out1 = [gather7_out1 as i64];
let concat5_out1: [i64; 3usize] = [
&unsqueeze18_out1[..],
&unsqueeze19_out1[..],
&unsqueeze20_out1[..],
]
.concat()
.try_into()
.unwrap();
let reshape5_out1 = concat5_out1;
let constant952_out1 = -1i64;
let mul11_out1 = {
let mut result = reshape5_out1;
let scalar = constant952_out1;
for result_item in result.iter_mut() {
*result_item = result_item.saturating_mul(scalar);
}
result
};
let max5_out1 = {
let __lhs = reshape5_out1;
let __rhs = mul11_out1;
core::array::from_fn::<i64, 3usize, _>(|__i| (__lhs[__i]).max(__rhs[__i]))
};
let expand5_out1 = {
let onnx_shape: [i64; 3usize] = max5_out1;
let input_dims = unsqueeze17_out1.dims();
let mut shape = onnx_shape;
#[allow(clippy::needless_range_loop)]
for i in 0..3usize {
let dim_offset = i;
if shape[dim_offset] == 1 && input_dims[i] > 1 {
shape[dim_offset] = input_dims[i] as i64;
}
}
unsqueeze17_out1.expand(shape)
};
let cast9_out1 = expand5_out1.float().cast(burn::tensor::DType::F32);
let gather8_out1 = shape8_out1[1] as i64;
let unsqueeze21_out1: Tensor<4> = cast9_out1.unsqueeze_dims::<4>(&[1]);
let unsqueeze22_out1 = [gather5_out1 as i64];
let unsqueeze23_out1 = [gather8_out1 as i64];
let unsqueeze24_out1 = [gather6_out1 as i64];
let unsqueeze25_out1 = [gather7_out1 as i64];
let concat6_out1: [i64; 4usize] = [
&unsqueeze22_out1[..],
&unsqueeze23_out1[..],
&unsqueeze24_out1[..],
&unsqueeze25_out1[..],
]
.concat()
.try_into()
.unwrap();
let reshape6_out1 = concat6_out1;
let constant961_out1 = -1i64;
let mul13_out1 = {
let mut result = reshape6_out1;
let scalar = constant961_out1;
for result_item in result.iter_mut() {
*result_item = result_item.saturating_mul(scalar);
}
result
};
let max6_out1 = {
let __lhs = reshape6_out1;
let __rhs = mul13_out1;
core::array::from_fn::<i64, 4usize, _>(|__i| (__lhs[__i]).max(__rhs[__i]))
};
let expand6_out1 = {
let onnx_shape: [i64; 4usize] = max6_out1;
let input_dims = unsqueeze21_out1.clone().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;
}
}
unsqueeze21_out1.clone().expand(shape)
};
let mul14_out1 = conv2d1_out1.mul(expand6_out1);
let relu1_out1 = burn::tensor::activation::relu(mul14_out1);
let shape15_out1: [i64; 4] = {
let axes = &relu1_out1.clone().dims()[0..4];
let mut output = [0i64; 4];
for i in 0..4 {
output[i] = axes[i] as i64;
}
output
};
let gather9_out1 = shape15_out1[0] as i64;
let gather10_out1 = shape15_out1[1] as i64;
let gather11_out1 = shape15_out1[2] as i64;
let gather12_out1 = shape15_out1[3] as i64;
let unsqueeze26_out1 = [gather9_out1 as i64];
let unsqueeze27_out1 = [gather10_out1 as i64];
let unsqueeze28_out1 = [gather11_out1 as i64];
let unsqueeze29_out1 = [gather12_out1 as i64];
let concat7_out1: [i64; 4usize] = [
&unsqueeze26_out1[..],
&unsqueeze27_out1[..],
&unsqueeze28_out1[..],
&unsqueeze29_out1[..],
]
.concat()
.try_into()
.unwrap();
let reshape7_out1 = concat7_out1;
let constant972_out1 = -1i64;
let mul16_out1 = {
let mut result = reshape7_out1;
let scalar = constant972_out1;
for result_item in result.iter_mut() {
*result_item = result_item.saturating_mul(scalar);
}
result
};
let max7_out1 = {
let __lhs = reshape7_out1;
let __rhs = mul16_out1;
core::array::from_fn::<i64, 4usize, _>(|__i| (__lhs[__i]).max(__rhs[__i]))
};
let expand7_out1 = {
let onnx_shape: [i64; 4usize] = max7_out1;
let input_dims = unsqueeze21_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;
}
}
unsqueeze21_out1.expand(shape)
};
let mul17_out1 = relu1_out1.mul(expand7_out1);
(mul17_out1, add3_out1)
}
}
#[derive(Module, Debug)]
pub struct Submodule2 {
conv2d2: Conv2d,
constant973: burn::module::Param<Tensor<1>>,
constant974: burn::module::Param<Tensor<1>>,
constant975: burn::module::Param<Tensor<1>>,
constant976: burn::module::Param<Tensor<1>>,
constant977: burn::module::Param<Tensor<1>>,
conv2d3: Conv2d,
#[module(skip)]
device: Device,
}
impl Submodule2 {
#[allow(unused_variables)]
pub fn new(device: &Device) -> Self {
let conv2d2 = Conv2dConfig::new([256, 256], [3, 3])
.with_stride([2, 2])
.with_padding(PaddingConfig2d::Explicit(1, 1, 1, 1))
.with_dilation([1, 1])
.with_groups(256)
.with_bias(true)
.init(device);
let constant973: 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 constant974: 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 constant975: 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([3f64]),
(device, burn::tensor::DType::F32),
),
device.clone(),
false,
[1].into(),
);
let constant976: 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([2f64]),
(device, burn::tensor::DType::F32),
),
device.clone(),
false,
[1].into(),
);
let constant977: 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 conv2d3 = Conv2dConfig::new([256, 256], [1, 1])
.with_stride([1, 1])
.with_padding(PaddingConfig2d::Valid)
.with_dilation([1, 1])
.with_groups(1)
.with_bias(true)
.init(device);
Self {
conv2d2,
constant973,
constant974,
constant975,
constant976,
constant977,
conv2d3,
device: device.clone(),
}
}
#[allow(clippy::let_and_return, clippy::approx_constant)]
pub fn forward(
&self,
mul17_out1: Tensor<4>,
add3_out1: Tensor<1>,
) -> (Tensor<4>, Tensor<1>) {
let conv2d2_out1 = self.conv2d2.forward(mul17_out1);
let constant973_out1 = self.constant973.val();
let add4_out1 = add3_out1.add(constant973_out1);
let constant974_out1 = self.constant974.val();
let add5_out1 = add4_out1.add(constant974_out1);
let constant975_out1 = self.constant975.val();
let sub2_out1 = add5_out1.sub(constant975_out1);
let constant976_out1 = self.constant976.val();
let div2_out1 = sub2_out1.div(constant976_out1);
let cast10_out1 = div2_out1.int().cast(burn::tensor::DType::I64);
let cast11_out1 = cast10_out1.float().cast(burn::tensor::DType::F32);
let constant977_out1 = self.constant977.val();
let add6_out1 = cast11_out1.add(constant977_out1);
let cast12_out1 = add6_out1.clone().int().cast(burn::tensor::DType::I64);
let shape20_out1: [i64; 4] = {
let axes = &conv2d2_out1.clone().dims()[0..4];
let mut output = [0i64; 4];
for i in 0..4 {
output[i] = axes[i] as i64;
}
output
};
let gather13_out1 = shape20_out1[0] as i64;
let gather14_out1 = shape20_out1[2] as i64;
let gather15_out1 = shape20_out1[3] as i64;
let range3_out1 = {
let start = 0i64;
let delta = 1i64;
assert!(delta != 0);
let n = (((gather14_out1 as i128 - start as i128) as f64) / delta as f64)
.ceil()
.max(0.0) as i64;
Tensor::arange(0..n, &self.device)
.cast(burn::tensor::DType::I64)
.mul_scalar(delta)
.add_scalar(start)
};
let unsqueeze30_out1 = [gather13_out1 as i64];
let unsqueeze31_out1 = [gather14_out1 as i64];
let concat8_out1: [i64; 2usize] = [&unsqueeze30_out1[..], &unsqueeze31_out1[..]]
.concat()
.try_into()
.unwrap();
let reshape8_out1 = concat8_out1;
let constant987_out1 = -1i64;
let mul19_out1 = {
let mut result = reshape8_out1;
let scalar = constant987_out1;
for result_item in result.iter_mut() {
*result_item = result_item.saturating_mul(scalar);
}
result
};
let max8_out1 = {
let __lhs = reshape8_out1;
let __rhs = mul19_out1;
core::array::from_fn::<i64, 2usize, _>(|__i| (__lhs[__i]).max(__rhs[__i]))
};
let expand8_out1 = {
let onnx_shape: [i64; 2usize] = max8_out1;
let input_dims = range3_out1.dims();
let mut shape = onnx_shape;
#[allow(clippy::needless_range_loop)]
for i in 0..1usize {
let dim_offset = 1usize + i;
if shape[dim_offset] == 1 && input_dims[i] > 1 {
shape[dim_offset] = input_dims[i] as i64;
}
}
range3_out1.expand(shape)
};
let unsqueeze32_out1: Tensor<2, Int> = cast12_out1.unsqueeze_dims::<2>(&[1]);
let less3_out1 = expand8_out1.lower(unsqueeze32_out1);
let unsqueeze33_out1: Tensor<3, Bool> = less3_out1.unsqueeze_dims::<3>(&[-1]);
let unsqueeze34_out1 = [gather13_out1 as i64];
let unsqueeze35_out1 = [gather14_out1 as i64];
let unsqueeze36_out1 = [gather15_out1 as i64];
let concat9_out1: [i64; 3usize] = [
&unsqueeze34_out1[..],
&unsqueeze35_out1[..],
&unsqueeze36_out1[..],
]
.concat()
.try_into()
.unwrap();
let reshape9_out1 = concat9_out1;
let constant995_out1 = -1i64;
let mul21_out1 = {
let mut result = reshape9_out1;
let scalar = constant995_out1;
for result_item in result.iter_mut() {
*result_item = result_item.saturating_mul(scalar);
}
result
};
let max9_out1 = {
let __lhs = reshape9_out1;
let __rhs = mul21_out1;
core::array::from_fn::<i64, 3usize, _>(|__i| (__lhs[__i]).max(__rhs[__i]))
};
let expand9_out1 = {
let onnx_shape: [i64; 3usize] = max9_out1;
let input_dims = unsqueeze33_out1.dims();
let mut shape = onnx_shape;
#[allow(clippy::needless_range_loop)]
for i in 0..3usize {
let dim_offset = i;
if shape[dim_offset] == 1 && input_dims[i] > 1 {
shape[dim_offset] = input_dims[i] as i64;
}
}
unsqueeze33_out1.expand(shape)
};
let cast14_out1 = expand9_out1.float().cast(burn::tensor::DType::F32);
let gather16_out1 = shape20_out1[1] as i64;
let unsqueeze37_out1: Tensor<4> = cast14_out1.unsqueeze_dims::<4>(&[1]);
let unsqueeze38_out1 = [gather13_out1 as i64];
let unsqueeze39_out1 = [gather16_out1 as i64];
let unsqueeze40_out1 = [gather14_out1 as i64];
let unsqueeze41_out1 = [gather15_out1 as i64];
let concat10_out1: [i64; 4usize] = [
&unsqueeze38_out1[..],
&unsqueeze39_out1[..],
&unsqueeze40_out1[..],
&unsqueeze41_out1[..],
]
.concat()
.try_into()
.unwrap();
let reshape10_out1 = concat10_out1;
let constant1004_out1 = -1i64;
let mul23_out1 = {
let mut result = reshape10_out1;
let scalar = constant1004_out1;
for result_item in result.iter_mut() {
*result_item = result_item.saturating_mul(scalar);
}
result
};
let max10_out1 = {
let __lhs = reshape10_out1;
let __rhs = mul23_out1;
core::array::from_fn::<i64, 4usize, _>(|__i| (__lhs[__i]).max(__rhs[__i]))
};
let expand10_out1 = {
let onnx_shape: [i64; 4usize] = max10_out1;
let input_dims = unsqueeze37_out1.clone().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;
}
}
unsqueeze37_out1.clone().expand(shape)
};
let mul24_out1 = conv2d2_out1.mul(expand10_out1);
let conv2d3_out1 = self.conv2d3.forward(mul24_out1);
let shape27_out1: [i64; 4] = {
let axes = &conv2d3_out1.clone().dims()[0..4];
let mut output = [0i64; 4];
for i in 0..4 {
output[i] = axes[i] as i64;
}
output
};
let gather17_out1 = shape27_out1[0] as i64;
let gather18_out1 = shape27_out1[1] as i64;
let gather19_out1 = shape27_out1[2] as i64;
let gather20_out1 = shape27_out1[3] as i64;
let unsqueeze42_out1 = [gather17_out1 as i64];
let unsqueeze43_out1 = [gather18_out1 as i64];
let unsqueeze44_out1 = [gather19_out1 as i64];
let unsqueeze45_out1 = [gather20_out1 as i64];
let concat11_out1: [i64; 4usize] = [
&unsqueeze42_out1[..],
&unsqueeze43_out1[..],
&unsqueeze44_out1[..],
&unsqueeze45_out1[..],
]
.concat()
.try_into()
.unwrap();
let reshape11_out1 = concat11_out1;
let constant1015_out1 = -1i64;
let mul26_out1 = {
let mut result = reshape11_out1;
let scalar = constant1015_out1;
for result_item in result.iter_mut() {
*result_item = result_item.saturating_mul(scalar);
}
result
};
let max11_out1 = {
let __lhs = reshape11_out1;
let __rhs = mul26_out1;
core::array::from_fn::<i64, 4usize, _>(|__i| (__lhs[__i]).max(__rhs[__i]))
};
let expand11_out1 = {
let onnx_shape: [i64; 4usize] = max11_out1;
let input_dims = unsqueeze37_out1.clone().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;
}
}
unsqueeze37_out1.clone().expand(shape)
};
let mul27_out1 = conv2d3_out1.mul(expand11_out1);
let relu2_out1 = burn::tensor::activation::relu(mul27_out1);
let shape32_out1: [i64; 4] = {
let axes = &relu2_out1.clone().dims()[0..4];
let mut output = [0i64; 4];
for i in 0..4 {
output[i] = axes[i] as i64;
}
output
};
let gather21_out1 = shape32_out1[0] as i64;
let gather22_out1 = shape32_out1[1] as i64;
let gather23_out1 = shape32_out1[2] as i64;
let gather24_out1 = shape32_out1[3] as i64;
let unsqueeze46_out1 = [gather21_out1 as i64];
let unsqueeze47_out1 = [gather22_out1 as i64];
let unsqueeze48_out1 = [gather23_out1 as i64];
let unsqueeze49_out1 = [gather24_out1 as i64];
let concat12_out1: [i64; 4usize] = [
&unsqueeze46_out1[..],
&unsqueeze47_out1[..],
&unsqueeze48_out1[..],
&unsqueeze49_out1[..],
]
.concat()
.try_into()
.unwrap();
let reshape12_out1 = concat12_out1;
let constant1026_out1 = -1i64;
let mul29_out1 = {
let mut result = reshape12_out1;
let scalar = constant1026_out1;
for result_item in result.iter_mut() {
*result_item = result_item.saturating_mul(scalar);
}
result
};
let max12_out1 = {
let __lhs = reshape12_out1;
let __rhs = mul29_out1;
core::array::from_fn::<i64, 4usize, _>(|__i| (__lhs[__i]).max(__rhs[__i]))
};
let expand12_out1 = {
let onnx_shape: [i64; 4usize] = max12_out1;
let input_dims = unsqueeze37_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;
}
}
unsqueeze37_out1.expand(shape)
};
let mul30_out1 = relu2_out1.mul(expand12_out1);
(mul30_out1, add6_out1)
}
}
#[derive(Module, Debug)]
pub struct Submodule3 {
conv2d4: Conv2d,
constant1027: burn::module::Param<Tensor<1>>,
constant1028: burn::module::Param<Tensor<1>>,
constant1029: burn::module::Param<Tensor<1>>,
constant1030: burn::module::Param<Tensor<1>>,
constant1031: burn::module::Param<Tensor<1>>,
conv2d5: Conv2d,
#[module(skip)]
device: Device,
}
impl Submodule3 {
#[allow(unused_variables)]
pub fn new(device: &Device) -> Self {
let conv2d4 = Conv2dConfig::new([256, 256], [3, 3])
.with_stride([2, 2])
.with_padding(PaddingConfig2d::Explicit(1, 1, 1, 1))
.with_dilation([1, 1])
.with_groups(256)
.with_bias(true)
.init(device);
let constant1027: 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 constant1028: 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 constant1029: 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([3f64]),
(device, burn::tensor::DType::F32),
),
device.clone(),
false,
[1].into(),
);
let constant1030: 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([2f64]),
(device, burn::tensor::DType::F32),
),
device.clone(),
false,
[1].into(),
);
let constant1031: 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 conv2d5 = Conv2dConfig::new([256, 256], [1, 1])
.with_stride([1, 1])
.with_padding(PaddingConfig2d::Valid)
.with_dilation([1, 1])
.with_groups(1)
.with_bias(true)
.init(device);
Self {
conv2d4,
constant1027,
constant1028,
constant1029,
constant1030,
constant1031,
conv2d5,
device: device.clone(),
}
}
#[allow(clippy::let_and_return, clippy::approx_constant)]
pub fn forward(
&self,
mul30_out1: Tensor<4>,
add6_out1: Tensor<1>,
) -> (Tensor<1>, Tensor<4>, Tensor<1, Int>) {
let conv2d4_out1 = self.conv2d4.forward(mul30_out1);
let constant1027_out1 = self.constant1027.val();
let add7_out1 = add6_out1.add(constant1027_out1);
let constant1028_out1 = self.constant1028.val();
let add8_out1 = add7_out1.add(constant1028_out1);
let constant1029_out1 = self.constant1029.val();
let sub3_out1 = add8_out1.sub(constant1029_out1);
let constant1030_out1 = self.constant1030.val();
let div3_out1 = sub3_out1.div(constant1030_out1);
let cast15_out1 = div3_out1.int().cast(burn::tensor::DType::I64);
let cast16_out1 = cast15_out1.float().cast(burn::tensor::DType::F32);
let constant1031_out1 = self.constant1031.val();
let add9_out1 = cast16_out1.add(constant1031_out1);
let cast17_out1 = add9_out1.clone().int().cast(burn::tensor::DType::I64);
let shape37_out1: [i64; 4] = {
let axes = &conv2d4_out1.clone().dims()[0..4];
let mut output = [0i64; 4];
for i in 0..4 {
output[i] = axes[i] as i64;
}
output
};
let gather25_out1 = shape37_out1[0] as i64;
let gather26_out1 = shape37_out1[2] as i64;
let gather27_out1 = shape37_out1[3] as i64;
let range4_out1 = {
let start = 0i64;
let delta = 1i64;
assert!(delta != 0);
let n = (((gather26_out1 as i128 - start as i128) as f64) / delta as f64)
.ceil()
.max(0.0) as i64;
Tensor::arange(0..n, &self.device)
.cast(burn::tensor::DType::I64)
.mul_scalar(delta)
.add_scalar(start)
};
let unsqueeze50_out1 = [gather25_out1 as i64];
let unsqueeze51_out1 = [gather26_out1 as i64];
let concat13_out1: [i64; 2usize] = [&unsqueeze50_out1[..], &unsqueeze51_out1[..]]
.concat()
.try_into()
.unwrap();
let reshape13_out1 = concat13_out1;
let constant1041_out1 = -1i64;
let mul32_out1 = {
let mut result = reshape13_out1;
let scalar = constant1041_out1;
for result_item in result.iter_mut() {
*result_item = result_item.saturating_mul(scalar);
}
result
};
let max13_out1 = {
let __lhs = reshape13_out1;
let __rhs = mul32_out1;
core::array::from_fn::<i64, 2usize, _>(|__i| (__lhs[__i]).max(__rhs[__i]))
};
let expand13_out1 = {
let onnx_shape: [i64; 2usize] = max13_out1;
let input_dims = range4_out1.dims();
let mut shape = onnx_shape;
#[allow(clippy::needless_range_loop)]
for i in 0..1usize {
let dim_offset = 1usize + i;
if shape[dim_offset] == 1 && input_dims[i] > 1 {
shape[dim_offset] = input_dims[i] as i64;
}
}
range4_out1.expand(shape)
};
let unsqueeze52_out1: Tensor<2, Int> = cast17_out1
.clone()
.unsqueeze_dims::<2>(&[1]);
let less4_out1 = expand13_out1.lower(unsqueeze52_out1);
let unsqueeze53_out1: Tensor<3, Bool> = less4_out1.unsqueeze_dims::<3>(&[-1]);
let unsqueeze54_out1 = [gather25_out1 as i64];
let unsqueeze55_out1 = [gather26_out1 as i64];
let unsqueeze56_out1 = [gather27_out1 as i64];
let concat14_out1: [i64; 3usize] = [
&unsqueeze54_out1[..],
&unsqueeze55_out1[..],
&unsqueeze56_out1[..],
]
.concat()
.try_into()
.unwrap();
let reshape14_out1 = concat14_out1;
let constant1049_out1 = -1i64;
let mul34_out1 = {
let mut result = reshape14_out1;
let scalar = constant1049_out1;
for result_item in result.iter_mut() {
*result_item = result_item.saturating_mul(scalar);
}
result
};
let max14_out1 = {
let __lhs = reshape14_out1;
let __rhs = mul34_out1;
core::array::from_fn::<i64, 3usize, _>(|__i| (__lhs[__i]).max(__rhs[__i]))
};
let expand14_out1 = {
let onnx_shape: [i64; 3usize] = max14_out1;
let input_dims = unsqueeze53_out1.dims();
let mut shape = onnx_shape;
#[allow(clippy::needless_range_loop)]
for i in 0..3usize {
let dim_offset = i;
if shape[dim_offset] == 1 && input_dims[i] > 1 {
shape[dim_offset] = input_dims[i] as i64;
}
}
unsqueeze53_out1.expand(shape)
};
let cast19_out1 = expand14_out1.float().cast(burn::tensor::DType::F32);
let gather28_out1 = shape37_out1[1] as i64;
let unsqueeze57_out1: Tensor<4> = cast19_out1.unsqueeze_dims::<4>(&[1]);
let unsqueeze58_out1 = [gather25_out1 as i64];
let unsqueeze59_out1 = [gather28_out1 as i64];
let unsqueeze60_out1 = [gather26_out1 as i64];
let unsqueeze61_out1 = [gather27_out1 as i64];
let concat15_out1: [i64; 4usize] = [
&unsqueeze58_out1[..],
&unsqueeze59_out1[..],
&unsqueeze60_out1[..],
&unsqueeze61_out1[..],
]
.concat()
.try_into()
.unwrap();
let reshape15_out1 = concat15_out1;
let constant1058_out1 = -1i64;
let mul36_out1 = {
let mut result = reshape15_out1;
let scalar = constant1058_out1;
for result_item in result.iter_mut() {
*result_item = result_item.saturating_mul(scalar);
}
result
};
let max15_out1 = {
let __lhs = reshape15_out1;
let __rhs = mul36_out1;
core::array::from_fn::<i64, 4usize, _>(|__i| (__lhs[__i]).max(__rhs[__i]))
};
let expand15_out1 = {
let onnx_shape: [i64; 4usize] = max15_out1;
let input_dims = unsqueeze57_out1.clone().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;
}
}
unsqueeze57_out1.clone().expand(shape)
};
let mul37_out1 = conv2d4_out1.mul(expand15_out1);
let conv2d5_out1 = self.conv2d5.forward(mul37_out1);
let shape44_out1: [i64; 4] = {
let axes = &conv2d5_out1.clone().dims()[0..4];
let mut output = [0i64; 4];
for i in 0..4 {
output[i] = axes[i] as i64;
}
output
};
let gather29_out1 = shape44_out1[0] as i64;
let gather30_out1 = shape44_out1[1] as i64;
let gather31_out1 = shape44_out1[2] as i64;
let gather32_out1 = shape44_out1[3] as i64;
let unsqueeze62_out1 = [gather29_out1 as i64];
let unsqueeze63_out1 = [gather30_out1 as i64];
let unsqueeze64_out1 = [gather31_out1 as i64];
let unsqueeze65_out1 = [gather32_out1 as i64];
let concat16_out1: [i64; 4usize] = [
&unsqueeze62_out1[..],
&unsqueeze63_out1[..],
&unsqueeze64_out1[..],
&unsqueeze65_out1[..],
]
.concat()
.try_into()
.unwrap();
let reshape16_out1 = concat16_out1;
let constant1069_out1 = -1i64;
let mul39_out1 = {
let mut result = reshape16_out1;
let scalar = constant1069_out1;
for result_item in result.iter_mut() {
*result_item = result_item.saturating_mul(scalar);
}
result
};
let max16_out1 = {
let __lhs = reshape16_out1;
let __rhs = mul39_out1;
core::array::from_fn::<i64, 4usize, _>(|__i| (__lhs[__i]).max(__rhs[__i]))
};
let expand16_out1 = {
let onnx_shape: [i64; 4usize] = max16_out1;
let input_dims = unsqueeze57_out1.clone().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;
}
}
unsqueeze57_out1.clone().expand(shape)
};
let mul40_out1 = conv2d5_out1.mul(expand16_out1);
let relu3_out1 = burn::tensor::activation::relu(mul40_out1);
let shape49_out1: [i64; 4] = {
let axes = &relu3_out1.clone().dims()[0..4];
let mut output = [0i64; 4];
for i in 0..4 {
output[i] = axes[i] as i64;
}
output
};
let gather33_out1 = shape49_out1[0] as i64;
let gather34_out1 = shape49_out1[1] as i64;
let gather35_out1 = shape49_out1[2] as i64;
let gather36_out1 = shape49_out1[3] as i64;
let unsqueeze66_out1 = [gather33_out1 as i64];
let unsqueeze67_out1 = [gather34_out1 as i64];
let unsqueeze68_out1 = [gather35_out1 as i64];
let unsqueeze69_out1 = [gather36_out1 as i64];
let concat17_out1: [i64; 4usize] = [
&unsqueeze66_out1[..],
&unsqueeze67_out1[..],
&unsqueeze68_out1[..],
&unsqueeze69_out1[..],
]
.concat()
.try_into()
.unwrap();
let reshape17_out1 = concat17_out1;
let constant1080_out1 = -1i64;
let mul42_out1 = {
let mut result = reshape17_out1;
let scalar = constant1080_out1;
for result_item in result.iter_mut() {
*result_item = result_item.saturating_mul(scalar);
}
result
};
let max17_out1 = {
let __lhs = reshape17_out1;
let __rhs = mul42_out1;
core::array::from_fn::<i64, 4usize, _>(|__i| (__lhs[__i]).max(__rhs[__i]))
};
let expand17_out1 = {
let onnx_shape: [i64; 4usize] = max17_out1;
let input_dims = unsqueeze57_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;
}
}
unsqueeze57_out1.expand(shape)
};
let mul43_out1 = relu3_out1.mul(expand17_out1);
(add9_out1, mul43_out1, cast17_out1)
}
}
#[derive(Module, Debug)]
pub struct Submodule4 {
linear1: Linear,
constant1087: burn::module::Param<Tensor<1>>,
constant1092: burn::module::Param<Tensor<3>>,
layernormalization1: LayerNorm,
linear2: Linear,
linear3: Linear,
constant1124: burn::module::Param<Tensor<1>>,
#[module(skip)]
device: Device,
}
impl Submodule4 {
#[allow(unused_variables)]
pub fn new(device: &Device) -> Self {
let linear1 = LinearConfig::new(4096, 1024).with_bias(true).init(device);
let constant1087: 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([32f64]),
(device, burn::tensor::DType::F32),
),
device.clone(),
false,
[1].into(),
);
let constant1092: burn::module::Param<Tensor<3>> = burn::module::Param::uninitialized(
burn::module::ParamId::new(),
move |device, _require_grad| Tensor::<
3,
>::zeros([1, 9999, 1024], (device, burn::tensor::DType::F32)),
device.clone(),
false,
[1, 9999, 1024].into(),
);
let layernormalization1 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let linear2 = LinearConfig::new(1024, 4096).with_bias(true).init(device);
let linear3 = LinearConfig::new(4096, 1024).with_bias(true).init(device);
let constant1124: 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([0.5f64]),
(device, burn::tensor::DType::F32),
),
device.clone(),
false,
[1].into(),
);
Self {
linear1,
constant1087,
constant1092,
layernormalization1,
linear2,
linear3,
constant1124,
device: device.clone(),
}
}
#[allow(clippy::let_and_return, clippy::approx_constant)]
pub fn forward(
&self,
add9_out1: Tensor<1>,
mul43_out1: Tensor<4>,
cast17_out1: Tensor<1, Int>,
) -> (Tensor<3>, Tensor<3>, Tensor<3, Bool>, Tensor<2, Bool>, Tensor<1, Int>) {
let cast20_out1 = add9_out1.int().cast(burn::tensor::DType::I64);
let shape54_out1: [i64; 4] = {
let axes = &mul43_out1.clone().dims()[0..4];
let mut output = [0i64; 4];
for i in 0..4 {
output[i] = axes[i] as i64;
}
output
};
let gather37_out1 = shape54_out1[0] as i64;
let gather38_out1 = shape54_out1[2] as i64;
let transpose2_out1 = mul43_out1.permute([0, 2, 1, 3]);
let unsqueeze70_out1 = [gather37_out1 as i64];
let unsqueeze71_out1 = [gather38_out1 as i64];
let constant1085_out1: [i64; 1] = [-1i64];
let concat18_out1: [i64; 3usize] = [
&unsqueeze70_out1[..],
&unsqueeze71_out1[..],
&constant1085_out1[..],
]
.concat()
.try_into()
.unwrap();
let reshape18_out1 = transpose2_out1.reshape(concat18_out1);
let linear1_out1 = self.linear1.forward(reshape18_out1);
let shape56_out1: [i64; 3] = {
let axes = &linear1_out1.clone().dims()[0..3];
let mut output = [0i64; 3];
for i in 0..3 {
output[i] = axes[i] as i64;
}
output
};
let gather39_out1 = shape56_out1[1] as i64;
let constant1087_out1 = self.constant1087.val();
let mul44_out1 = linear1_out1
.mul((constant1087_out1).unsqueeze_dims(&[0isize, 1isize]));
let shape57_out1: [i64; 3] = {
let axes = &mul44_out1.clone().dims()[0..3];
let mut output = [0i64; 3];
for i in 0..3 {
output[i] = axes[i] as i64;
}
output
};
let gather40_out1 = shape57_out1[1] as i64;
let constant1089_out1 = 5000i64;
let sub4_out1 = constant1089_out1 - gather40_out1;
let constant1090_out1 = 5000i64;
let add10_out1 = constant1090_out1 + gather40_out1;
let constant1091_out1 = 1i64;
let sub5_out1 = add10_out1 - constant1091_out1;
let unsqueeze72_out1 = [sub4_out1 as i64];
let unsqueeze73_out1 = [sub5_out1 as i64];
let constant1092_out1 = self.constant1092.val();
let slice1_out1 = constant1092_out1
.slice(s![.., unsqueeze72_out1[0]..unsqueeze73_out1[0], ..]);
let unsqueeze74_out1 = [gather39_out1 as i64];
let unsqueeze75_out1 = [gather39_out1 as i64];
let constant1097_out1: [i64; 1] = [1i64];
let concat19_out1: [i64; 3usize] = [
&constant1097_out1[..],
&unsqueeze74_out1[..],
&unsqueeze75_out1[..],
]
.concat()
.try_into()
.unwrap();
let constantofshape18_out1 = Tensor::<3, Int>::ones(concat19_out1, &self.device)
.bool();
let range5_out1 = {
let start = 0i64;
let delta = 1i64;
assert!(delta != 0);
let n = (((gather39_out1 as i128 - start as i128) as f64) / delta as f64)
.ceil()
.max(0.0) as i64;
Tensor::arange(0..n, &self.device)
.cast(burn::tensor::DType::I64)
.mul_scalar(delta)
.add_scalar(start)
};
let shape58_out1: [i64; 1] = {
let axes = &cast17_out1.clone().dims()[0..1];
let mut output = [0i64; 1];
for i in 0..1 {
output[i] = axes[i] as i64;
}
output
};
let gather41_out1 = shape58_out1[0] as i64;
let unsqueeze76_out1 = [gather41_out1 as i64];
let constant1104_out1: [i64; 1] = [-1i64];
let concat20_out1: [i64; 2usize] = [
&unsqueeze76_out1[..],
&constant1104_out1[..],
]
.concat()
.try_into()
.unwrap();
let reshape19_out1 = concat20_out1;
let constant1107_out1 = -1i64;
let mul46_out1 = {
let mut result = reshape19_out1;
let scalar = constant1107_out1;
for result_item in result.iter_mut() {
*result_item = result_item.saturating_mul(scalar);
}
result
};
let max18_out1 = {
let __lhs = reshape19_out1;
let __rhs = mul46_out1;
core::array::from_fn::<i64, 2usize, _>(|__i| (__lhs[__i]).max(__rhs[__i]))
};
let expand18_out1 = {
let onnx_shape: [i64; 2usize] = max18_out1;
let input_dims = range5_out1.dims();
let mut shape = onnx_shape;
#[allow(clippy::needless_range_loop)]
for i in 0..1usize {
let dim_offset = 1usize + i;
if shape[dim_offset] == 1 && input_dims[i] > 1 {
shape[dim_offset] = input_dims[i] as i64;
}
}
range5_out1.expand(shape)
};
let unsqueeze77_out1: Tensor<2, Int> = cast17_out1.unsqueeze_dims::<2>(&[-1]);
let less5_out1 = expand18_out1.lower(unsqueeze77_out1);
let unsqueeze78_out1: Tensor<3, Bool> = less5_out1
.clone()
.unsqueeze_dims::<3>(&[1]);
let unsqueeze80_out1 = [gather39_out1 as i64];
let constant1115_out1: [i64; 1] = [1i64];
let constant1113_out1: [i64; 1] = [1i64];
let concat22_out1: [i64; 3usize] = [
&constant1113_out1[..],
&unsqueeze80_out1[..],
&constant1115_out1[..],
]
.concat()
.try_into()
.unwrap();
let tile1_out1 = {
let repeats: alloc::vec::Vec<usize> = concat22_out1
.iter()
.map(|&v| v as usize)
.collect();
unsqueeze78_out1.repeat(&repeats)
};
let transpose3_out1 = tile1_out1.clone().permute([0, 2, 1]);
let and1_out1 = tile1_out1.bool_and(transpose3_out1);
let unsqueeze81_out1 = [gather39_out1 as i64];
let slice2_out1 = constantofshape18_out1
.slice(s![.., 0..unsqueeze81_out1[0], ..]);
let unsqueeze82_out1 = [gather39_out1 as i64];
let slice3_out1 = slice2_out1.slice(s![.., .., 0..unsqueeze82_out1[0]]);
let and2_out1 = and1_out1.bool_and(slice3_out1);
let not1_out1 = and2_out1.bool_not();
let not2_out1 = less5_out1.bool_not();
let layernormalization1_out1 = {
let dtype = mul44_out1.clone().dtype();
self.layernormalization1
.forward(mul44_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let linear2_out1 = self.linear2.forward(layernormalization1_out1);
let sigmoid1_out1 = burn::tensor::activation::sigmoid(linear2_out1.clone());
let mul47_out1 = linear2_out1.mul(sigmoid1_out1);
let linear3_out1 = self.linear3.forward(mul47_out1);
let constant1124_out1 = self.constant1124.val();
let mul48_out1 = linear3_out1
.mul((constant1124_out1).unsqueeze_dims(&[0isize, 1isize]));
let add11_out1 = mul44_out1.add(mul48_out1);
(add11_out1, slice1_out1, not1_out1, not2_out1, cast20_out1)
}
}
#[derive(Module, Debug)]
pub struct Submodule5 {
layernormalization2: LayerNorm,
linear4: Linear,
linear5: Linear,
linear6: Linear,
linear7: Linear,
constant24: burn::module::Param<Tensor<2>>,
constant25: burn::module::Param<Tensor<2>>,
constant1267: burn::module::Param<Tensor<1>>,
linear8: Linear,
#[module(skip)]
device: Device,
}
impl Submodule5 {
#[allow(unused_variables)]
pub fn new(device: &Device) -> Self {
let layernormalization2 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.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(false).init(device);
let constant24: burn::module::Param<Tensor<2>> = burn::module::Param::uninitialized(
burn::module::ParamId::new(),
move |device, _require_grad| Tensor::<
2,
>::zeros([8, 128], (device, burn::tensor::DType::F32)),
device.clone(),
false,
[8, 128].into(),
);
let constant25: burn::module::Param<Tensor<2>> = burn::module::Param::uninitialized(
burn::module::ParamId::new(),
move |device, _require_grad| Tensor::<
2,
>::zeros([8, 128], (device, burn::tensor::DType::F32)),
device.clone(),
false,
[8, 128].into(),
);
let constant1267: 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([11.313708305358887f64]),
(device, burn::tensor::DType::F32),
),
device.clone(),
false,
[1].into(),
);
let linear8 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
Self {
layernormalization2,
linear4,
linear5,
linear6,
linear7,
constant24,
constant25,
constant1267,
linear8,
device: device.clone(),
}
}
#[allow(clippy::let_and_return, clippy::approx_constant)]
pub fn forward(
&self,
add11_out1: Tensor<3>,
slice1_out1: Tensor<3>,
not1_out1: Tensor<3, Bool>,
) -> (Tensor<3>, Tensor<4, Bool>) {
let layernormalization2_out1 = {
let dtype = add11_out1.clone().dtype();
self.layernormalization2
.forward(add11_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let shape61_out1: [i64; 3] = {
let axes = &layernormalization2_out1.clone().dims()[0..3];
let mut output = [0i64; 3];
for i in 0..3 {
output[i] = axes[i] as i64;
}
output
};
let gather42_out1 = shape61_out1[0] as i64;
let linear4_out1 = self.linear4.forward(layernormalization2_out1.clone());
let unsqueeze83_out1 = [gather42_out1 as i64];
let constant1129_out1: [i64; 1] = [128i64];
let constant1127_out1: [i64; 1] = [-1i64];
let constant1128_out1: [i64; 1] = [8i64];
let concat23_out1: [i64; 4usize] = [
&unsqueeze83_out1[..],
&constant1127_out1[..],
&constant1128_out1[..],
&constant1129_out1[..],
]
.concat()
.try_into()
.unwrap();
let unsqueeze84_out1 = [gather42_out1 as i64];
let constant1133_out1: [i64; 1] = [128i64];
let constant1131_out1: [i64; 1] = [-1i64];
let constant1132_out1: [i64; 1] = [8i64];
let concat24_out1: [i64; 4usize] = [
&unsqueeze84_out1[..],
&constant1131_out1[..],
&constant1132_out1[..],
&constant1133_out1[..],
]
.concat()
.try_into()
.unwrap();
let unsqueeze85_out1 = [gather42_out1 as i64];
let constant1137_out1: [i64; 1] = [128i64];
let constant1135_out1: [i64; 1] = [-1i64];
let constant1136_out1: [i64; 1] = [8i64];
let concat25_out1: [i64; 4usize] = [
&unsqueeze85_out1[..],
&constant1135_out1[..],
&constant1136_out1[..],
&constant1137_out1[..],
]
.concat()
.try_into()
.unwrap();
let reshape20_out1 = linear4_out1.reshape(concat23_out1);
let linear5_out1 = self.linear5.forward(layernormalization2_out1.clone());
let reshape21_out1 = linear5_out1.reshape(concat24_out1);
let linear6_out1 = self.linear6.forward(layernormalization2_out1);
let reshape22_out1 = linear6_out1.reshape(concat25_out1);
let transpose4_out1 = reshape22_out1.permute([0, 2, 1, 3]);
let linear7_out1 = self.linear7.forward(slice1_out1);
let reshape23_out1 = linear7_out1.reshape([1, -1, 8, 128]);
let constant24_out1 = self.constant24.val();
let add12_out1 = reshape20_out1
.clone()
.add((constant24_out1).unsqueeze_dims(&[0isize, 1isize]));
let transpose5_out1 = add12_out1.permute([0, 2, 1, 3]);
let constant25_out1 = self.constant25.val();
let add13_out1 = reshape20_out1
.add((constant25_out1).unsqueeze_dims(&[0isize, 1isize]));
let transpose6_out1 = add13_out1.permute([0, 2, 1, 3]);
let transpose7_out1 = reshape23_out1.permute([0, 2, 3, 1]);
let matmul8_out1 = transpose6_out1.matmul(transpose7_out1);
let shape63_out1: [i64; 4] = {
let axes = &matmul8_out1.clone().dims()[0..4];
let mut output = [0i64; 4];
for i in 0..4 {
output[i] = axes[i] as i64;
}
output
};
let gather44_out1 = shape63_out1[0] as i64;
let gather45_out1 = shape63_out1[1] as i64;
let gather46_out1 = shape63_out1[2] as i64;
let gather47_out1 = shape63_out1[3] as i64;
let pad1_out1 = matmul8_out1
.pad(
[(0usize, 0usize), (0usize, 0usize), (0usize, 0usize), (1usize, 0usize)],
burn::tensor::ops::PadMode::Constant(0f32),
);
let unsqueeze110_out1 = [gather44_out1 as i64];
let unsqueeze111_out1 = [gather45_out1 as i64];
let unsqueeze112_out1 = [gather46_out1 as i64];
let constant1249_out1: [i64; 1] = [-1i64];
let concat51_out1: [i64; 4usize] = [
&unsqueeze110_out1[..],
&unsqueeze111_out1[..],
&constant1249_out1[..],
&unsqueeze112_out1[..],
]
.concat()
.try_into()
.unwrap();
let reshape26_out1 = pad1_out1.reshape(concat51_out1);
let slice5_out1 = reshape26_out1.slice(s![.., .., 1.., ..]);
let unsqueeze113_out1 = [gather44_out1 as i64];
let unsqueeze114_out1 = [gather45_out1 as i64];
let unsqueeze115_out1 = [gather46_out1 as i64];
let unsqueeze116_out1 = [gather47_out1 as i64];
let concat52_out1: [i64; 4usize] = [
&unsqueeze113_out1[..],
&unsqueeze114_out1[..],
&unsqueeze115_out1[..],
&unsqueeze116_out1[..],
]
.concat()
.try_into()
.unwrap();
let reshape27_out1 = slice5_out1.reshape(concat52_out1);
let transpose9_out1 = reshape21_out1.permute([0, 2, 3, 1]);
let matmul9_out1 = transpose5_out1.matmul(transpose9_out1);
let shape67_out1: [i64; 4] = {
let axes = &matmul9_out1.clone().dims()[0..4];
let mut output = [0i64; 4];
for i in 0..4 {
output[i] = axes[i] as i64;
}
output
};
let slice6_out1: [i64; 1] = shape67_out1[3..4].try_into().unwrap();
let squeeze1_out1 = slice6_out1[0] as i64;
let unsqueeze117_out1 = [squeeze1_out1 as i64];
let slice7_out1 = reshape27_out1.slice(s![.., .., .., 0..unsqueeze117_out1[0]]);
let add14_out1 = matmul9_out1.add(slice7_out1);
let constant1267_out1 = self.constant1267.val();
let div4_out1 = add14_out1
.div((constant1267_out1).unsqueeze_dims(&[0isize, 1isize, 2isize]));
let shape68_out1: [i64; 4] = {
let axes = &transpose4_out1.clone().dims()[0..4];
let mut output = [0i64; 4];
for i in 0..4 {
output[i] = axes[i] as i64;
}
output
};
let gather48_out1 = shape68_out1[0] as i64;
let unsqueeze118_out1: Tensor<4, Bool> = not1_out1.unsqueeze_dims::<4>(&[1]);
let constant1270_out1 = -10000f32;
let where19_out1 = div4_out1
.mask_fill(unsqueeze118_out1.clone(), constant1270_out1);
let softmax1_out1 = burn::tensor::activation::softmax(where19_out1, 3);
let constant1271_out1 = 0f32;
let where20_out1 = softmax1_out1
.mask_fill(unsqueeze118_out1.clone(), constant1271_out1);
let matmul10_out1 = where20_out1.matmul(transpose4_out1);
let transpose10_out1 = matmul10_out1.permute([0, 2, 1, 3]);
let unsqueeze119_out1 = [gather48_out1 as i64];
let constant1274_out1: [i64; 1] = [1024i64];
let constant1273_out1: [i64; 1] = [-1i64];
let concat53_out1: [i64; 3usize] = [
&unsqueeze119_out1[..],
&constant1273_out1[..],
&constant1274_out1[..],
]
.concat()
.try_into()
.unwrap();
let reshape28_out1 = transpose10_out1.reshape(concat53_out1);
let linear8_out1 = self.linear8.forward(reshape28_out1);
let add15_out1 = add11_out1.add(linear8_out1);
(add15_out1, unsqueeze118_out1)
}
}
#[derive(Module, Debug)]
pub struct Submodule6 {
layernormalization3: LayerNorm,
conv1d1: Conv1d,
conv1d2: Conv1d,
conv1d3: Conv1d,
layernormalization4: LayerNorm,
linear9: Linear,
linear10: Linear,
constant1285: burn::module::Param<Tensor<1>>,
layernormalization5: LayerNorm,
layernormalization6: LayerNorm,
linear11: Linear,
linear12: Linear,
constant1286: burn::module::Param<Tensor<1>>,
layernormalization7: LayerNorm,
linear13: Linear,
linear14: Linear,
linear15: Linear,
linear16: Linear,
constant48: burn::module::Param<Tensor<2>>,
constant49: burn::module::Param<Tensor<2>>,
constant1332: burn::module::Param<Tensor<1>>,
linear17: Linear,
#[module(skip)]
device: Device,
}
impl Submodule6 {
#[allow(unused_variables)]
pub fn new(device: &Device) -> Self {
let layernormalization3 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let conv1d1 = Conv1dConfig::new(1024, 2048, 1)
.with_stride(1)
.with_padding(PaddingConfig1d::Valid)
.with_dilation(1)
.with_groups(1)
.with_bias(true)
.init(device);
let conv1d2 = Conv1dConfig::new(1024, 1024, 9)
.with_stride(1)
.with_padding(PaddingConfig1d::Valid)
.with_dilation(1)
.with_groups(1024)
.with_bias(true)
.init(device);
let conv1d3 = Conv1dConfig::new(1024, 1024, 1)
.with_stride(1)
.with_padding(PaddingConfig1d::Valid)
.with_dilation(1)
.with_groups(1)
.with_bias(true)
.init(device);
let layernormalization4 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let linear9 = LinearConfig::new(1024, 4096).with_bias(true).init(device);
let linear10 = LinearConfig::new(4096, 1024).with_bias(true).init(device);
let constant1285: 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([0.5f64]),
(device, burn::tensor::DType::F32),
),
device.clone(),
false,
[1].into(),
);
let layernormalization5 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let layernormalization6 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let linear11 = LinearConfig::new(1024, 4096).with_bias(true).init(device);
let linear12 = LinearConfig::new(4096, 1024).with_bias(true).init(device);
let constant1286: 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([0.5f64]),
(device, burn::tensor::DType::F32),
),
device.clone(),
false,
[1].into(),
);
let layernormalization7 = 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 linear15 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
let linear16 = LinearConfig::new(1024, 1024).with_bias(false).init(device);
let constant48: burn::module::Param<Tensor<2>> = burn::module::Param::uninitialized(
burn::module::ParamId::new(),
move |device, _require_grad| Tensor::<
2,
>::zeros([8, 128], (device, burn::tensor::DType::F32)),
device.clone(),
false,
[8, 128].into(),
);
let constant49: burn::module::Param<Tensor<2>> = burn::module::Param::uninitialized(
burn::module::ParamId::new(),
move |device, _require_grad| Tensor::<
2,
>::zeros([8, 128], (device, burn::tensor::DType::F32)),
device.clone(),
false,
[8, 128].into(),
);
let constant1332: 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([11.313708305358887f64]),
(device, burn::tensor::DType::F32),
),
device.clone(),
false,
[1].into(),
);
let linear17 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
Self {
layernormalization3,
conv1d1,
conv1d2,
conv1d3,
layernormalization4,
linear9,
linear10,
constant1285,
layernormalization5,
layernormalization6,
linear11,
linear12,
constant1286,
layernormalization7,
linear13,
linear14,
linear15,
linear16,
constant48,
constant49,
constant1332,
linear17,
device: device.clone(),
}
}
#[allow(clippy::let_and_return, clippy::approx_constant)]
pub fn forward(
&self,
add15_out1: Tensor<3>,
not2_out1: Tensor<2, Bool>,
slice1_out1: Tensor<3>,
unsqueeze118_out1: Tensor<4, Bool>,
) -> (Tensor<3>, Tensor<3, Bool>) {
let layernormalization3_out1 = {
let dtype = add15_out1.clone().dtype();
self.layernormalization3
.forward(add15_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let transpose11_out1 = layernormalization3_out1.permute([0, 2, 1]);
let conv1d1_out1 = self.conv1d1.forward(transpose11_out1);
let split_tensors = conv1d1_out1.split_with_sizes([1024, 1024].into(), 1);
let [split1_out1, split1_out2] = split_tensors.try_into().unwrap();
let sigmoid2_out1 = burn::tensor::activation::sigmoid(split1_out2);
let mul49_out1 = split1_out1.mul(sigmoid2_out1);
let unsqueeze120_out1: Tensor<3, Bool> = not2_out1.unsqueeze_dims::<3>(&[1]);
let constant1276_out1 = 0f32;
let where21_out1 = mul49_out1
.mask_fill(unsqueeze120_out1.clone(), constant1276_out1);
let pad2_out1 = where21_out1
.pad(
[(0usize, 0usize), (0usize, 0usize), (4usize, 4usize)],
burn::tensor::ops::PadMode::Constant(0f32),
);
let conv1d2_out1 = self.conv1d2.forward(pad2_out1);
let sigmoid3_out1 = burn::tensor::activation::sigmoid(conv1d2_out1.clone());
let mul50_out1 = conv1d2_out1.mul(sigmoid3_out1);
let conv1d3_out1 = self.conv1d3.forward(mul50_out1);
let transpose13_out1 = conv1d3_out1.permute([0, 2, 1]);
let add16_out1 = add15_out1.add(transpose13_out1);
let layernormalization4_out1 = {
let dtype = add16_out1.clone().dtype();
self.layernormalization4
.forward(add16_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let linear9_out1 = self.linear9.forward(layernormalization4_out1);
let sigmoid4_out1 = burn::tensor::activation::sigmoid(linear9_out1.clone());
let mul51_out1 = linear9_out1.mul(sigmoid4_out1);
let linear10_out1 = self.linear10.forward(mul51_out1);
let constant1285_out1 = self.constant1285.val();
let mul52_out1 = linear10_out1
.mul((constant1285_out1).unsqueeze_dims(&[0isize, 1isize]));
let add17_out1 = add16_out1.add(mul52_out1);
let layernormalization5_out1 = {
let dtype = add17_out1.dtype();
self.layernormalization5
.forward(add17_out1.cast(burn::tensor::DType::F32))
.cast(dtype)
};
let layernormalization6_out1 = {
let dtype = layernormalization5_out1.clone().dtype();
self.layernormalization6
.forward(layernormalization5_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let linear11_out1 = self.linear11.forward(layernormalization6_out1);
let sigmoid5_out1 = burn::tensor::activation::sigmoid(linear11_out1.clone());
let mul53_out1 = linear11_out1.mul(sigmoid5_out1);
let linear12_out1 = self.linear12.forward(mul53_out1);
let constant1286_out1 = self.constant1286.val();
let mul54_out1 = linear12_out1
.mul((constant1286_out1).unsqueeze_dims(&[0isize, 1isize]));
let add18_out1 = layernormalization5_out1.add(mul54_out1);
let layernormalization7_out1 = {
let dtype = add18_out1.clone().dtype();
self.layernormalization7
.forward(add18_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let shape69_out1: [i64; 3] = {
let axes = &layernormalization7_out1.clone().dims()[0..3];
let mut output = [0i64; 3];
for i in 0..3 {
output[i] = axes[i] as i64;
}
output
};
let gather49_out1 = shape69_out1[0] as i64;
let linear13_out1 = self.linear13.forward(layernormalization7_out1.clone());
let unsqueeze121_out1 = [gather49_out1 as i64];
let constant1291_out1: [i64; 1] = [128i64];
let constant1289_out1: [i64; 1] = [-1i64];
let constant1290_out1: [i64; 1] = [8i64];
let concat55_out1: [i64; 4usize] = [
&unsqueeze121_out1[..],
&constant1289_out1[..],
&constant1290_out1[..],
&constant1291_out1[..],
]
.concat()
.try_into()
.unwrap();
let unsqueeze122_out1 = [gather49_out1 as i64];
let constant1295_out1: [i64; 1] = [128i64];
let constant1293_out1: [i64; 1] = [-1i64];
let constant1294_out1: [i64; 1] = [8i64];
let concat56_out1: [i64; 4usize] = [
&unsqueeze122_out1[..],
&constant1293_out1[..],
&constant1294_out1[..],
&constant1295_out1[..],
]
.concat()
.try_into()
.unwrap();
let unsqueeze123_out1 = [gather49_out1 as i64];
let constant1299_out1: [i64; 1] = [128i64];
let constant1297_out1: [i64; 1] = [-1i64];
let constant1298_out1: [i64; 1] = [8i64];
let concat57_out1: [i64; 4usize] = [
&unsqueeze123_out1[..],
&constant1297_out1[..],
&constant1298_out1[..],
&constant1299_out1[..],
]
.concat()
.try_into()
.unwrap();
let reshape31_out1 = linear13_out1.reshape(concat55_out1);
let linear14_out1 = self.linear14.forward(layernormalization7_out1.clone());
let reshape32_out1 = linear14_out1.reshape(concat56_out1);
let linear15_out1 = self.linear15.forward(layernormalization7_out1);
let reshape33_out1 = linear15_out1.reshape(concat57_out1);
let transpose14_out1 = reshape33_out1.permute([0, 2, 1, 3]);
let linear16_out1 = self.linear16.forward(slice1_out1);
let reshape34_out1 = linear16_out1.reshape([1, -1, 8, 128]);
let constant48_out1 = self.constant48.val();
let add19_out1 = reshape31_out1
.clone()
.add((constant48_out1).unsqueeze_dims(&[0isize, 1isize]));
let transpose15_out1 = add19_out1.permute([0, 2, 1, 3]);
let constant49_out1 = self.constant49.val();
let add20_out1 = reshape31_out1
.add((constant49_out1).unsqueeze_dims(&[0isize, 1isize]));
let transpose16_out1 = add20_out1.permute([0, 2, 1, 3]);
let transpose17_out1 = reshape34_out1.permute([0, 2, 3, 1]);
let matmul20_out1 = transpose16_out1.matmul(transpose17_out1);
let shape70_out1: [i64; 4] = {
let axes = &matmul20_out1.clone().dims()[0..4];
let mut output = [0i64; 4];
for i in 0..4 {
output[i] = axes[i] as i64;
}
output
};
let gather50_out1 = shape70_out1[0] as i64;
let gather51_out1 = shape70_out1[1] as i64;
let gather52_out1 = shape70_out1[2] as i64;
let gather53_out1 = shape70_out1[3] as i64;
let pad3_out1 = matmul20_out1
.pad(
[(0usize, 0usize), (0usize, 0usize), (0usize, 0usize), (1usize, 0usize)],
burn::tensor::ops::PadMode::Constant(0f32),
);
let unsqueeze124_out1 = [gather50_out1 as i64];
let unsqueeze125_out1 = [gather51_out1 as i64];
let unsqueeze126_out1 = [gather52_out1 as i64];
let constant1314_out1: [i64; 1] = [-1i64];
let concat59_out1: [i64; 4usize] = [
&unsqueeze124_out1[..],
&unsqueeze125_out1[..],
&constant1314_out1[..],
&unsqueeze126_out1[..],
]
.concat()
.try_into()
.unwrap();
let reshape37_out1 = pad3_out1.reshape(concat59_out1);
let slice10_out1 = reshape37_out1.slice(s![.., .., 1.., ..]);
let unsqueeze127_out1 = [gather50_out1 as i64];
let unsqueeze128_out1 = [gather51_out1 as i64];
let unsqueeze129_out1 = [gather52_out1 as i64];
let unsqueeze130_out1 = [gather53_out1 as i64];
let concat60_out1: [i64; 4usize] = [
&unsqueeze127_out1[..],
&unsqueeze128_out1[..],
&unsqueeze129_out1[..],
&unsqueeze130_out1[..],
]
.concat()
.try_into()
.unwrap();
let reshape38_out1 = slice10_out1.reshape(concat60_out1);
let transpose19_out1 = reshape32_out1.permute([0, 2, 3, 1]);
let matmul21_out1 = transpose15_out1.matmul(transpose19_out1);
let shape74_out1: [i64; 4] = {
let axes = &matmul21_out1.clone().dims()[0..4];
let mut output = [0i64; 4];
for i in 0..4 {
output[i] = axes[i] as i64;
}
output
};
let slice11_out1: [i64; 1] = shape74_out1[3..4].try_into().unwrap();
let squeeze2_out1 = slice11_out1[0] as i64;
let unsqueeze131_out1 = [squeeze2_out1 as i64];
let slice12_out1 = reshape38_out1.slice(s![.., .., .., 0..unsqueeze131_out1[0]]);
let add21_out1 = matmul21_out1.add(slice12_out1);
let constant1332_out1 = self.constant1332.val();
let div5_out1 = add21_out1
.div((constant1332_out1).unsqueeze_dims(&[0isize, 1isize, 2isize]));
let shape75_out1: [i64; 4] = {
let axes = &transpose14_out1.clone().dims()[0..4];
let mut output = [0i64; 4];
for i in 0..4 {
output[i] = axes[i] as i64;
}
output
};
let gather54_out1 = shape75_out1[0] as i64;
let constant1334_out1 = -10000f32;
let where22_out1 = div5_out1
.mask_fill(unsqueeze118_out1.clone(), constant1334_out1);
let softmax2_out1 = burn::tensor::activation::softmax(where22_out1, 3);
let constant1335_out1 = 0f32;
let where23_out1 = softmax2_out1.mask_fill(unsqueeze118_out1, constant1335_out1);
let matmul22_out1 = where23_out1.matmul(transpose14_out1);
let transpose20_out1 = matmul22_out1.permute([0, 2, 1, 3]);
let unsqueeze132_out1 = [gather54_out1 as i64];
let constant1338_out1: [i64; 1] = [1024i64];
let constant1337_out1: [i64; 1] = [-1i64];
let concat61_out1: [i64; 3usize] = [
&unsqueeze132_out1[..],
&constant1337_out1[..],
&constant1338_out1[..],
]
.concat()
.try_into()
.unwrap();
let reshape39_out1 = transpose20_out1.reshape(concat61_out1);
let linear17_out1 = self.linear17.forward(reshape39_out1);
let add22_out1 = add18_out1.add(linear17_out1);
(add22_out1, unsqueeze120_out1)
}
}
#[derive(Module, Debug)]
pub struct Submodule7 {
layernormalization8: LayerNorm,
conv1d4: Conv1d,
conv1d5: Conv1d,
conv1d6: Conv1d,
layernormalization9: LayerNorm,
linear18: Linear,
linear19: Linear,
constant1348: burn::module::Param<Tensor<1>>,
layernormalization10: LayerNorm,
layernormalization11: LayerNorm,
linear20: Linear,
linear21: Linear,
constant1349: burn::module::Param<Tensor<1>>,
layernormalization12: LayerNorm,
linear22: Linear,
linear23: Linear,
linear24: Linear,
linear25: Linear,
constant72: burn::module::Param<Tensor<2>>,
constant73: burn::module::Param<Tensor<2>>,
constant1395: burn::module::Param<Tensor<1>>,
linear26: Linear,
#[module(skip)]
device: Device,
}
impl Submodule7 {
#[allow(unused_variables)]
pub fn new(device: &Device) -> Self {
let layernormalization8 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let conv1d4 = Conv1dConfig::new(1024, 2048, 1)
.with_stride(1)
.with_padding(PaddingConfig1d::Valid)
.with_dilation(1)
.with_groups(1)
.with_bias(true)
.init(device);
let conv1d5 = Conv1dConfig::new(1024, 1024, 9)
.with_stride(1)
.with_padding(PaddingConfig1d::Valid)
.with_dilation(1)
.with_groups(1024)
.with_bias(true)
.init(device);
let conv1d6 = Conv1dConfig::new(1024, 1024, 1)
.with_stride(1)
.with_padding(PaddingConfig1d::Valid)
.with_dilation(1)
.with_groups(1)
.with_bias(true)
.init(device);
let layernormalization9 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let linear18 = LinearConfig::new(1024, 4096).with_bias(true).init(device);
let linear19 = LinearConfig::new(4096, 1024).with_bias(true).init(device);
let constant1348: 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([0.5f64]),
(device, burn::tensor::DType::F32),
),
device.clone(),
false,
[1].into(),
);
let layernormalization10 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let layernormalization11 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let linear20 = LinearConfig::new(1024, 4096).with_bias(true).init(device);
let linear21 = LinearConfig::new(4096, 1024).with_bias(true).init(device);
let constant1349: 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([0.5f64]),
(device, burn::tensor::DType::F32),
),
device.clone(),
false,
[1].into(),
);
let layernormalization12 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let linear22 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
let linear23 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
let linear24 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
let linear25 = LinearConfig::new(1024, 1024).with_bias(false).init(device);
let constant72: burn::module::Param<Tensor<2>> = burn::module::Param::uninitialized(
burn::module::ParamId::new(),
move |device, _require_grad| Tensor::<
2,
>::zeros([8, 128], (device, burn::tensor::DType::F32)),
device.clone(),
false,
[8, 128].into(),
);
let constant73: burn::module::Param<Tensor<2>> = burn::module::Param::uninitialized(
burn::module::ParamId::new(),
move |device, _require_grad| Tensor::<
2,
>::zeros([8, 128], (device, burn::tensor::DType::F32)),
device.clone(),
false,
[8, 128].into(),
);
let constant1395: 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([11.313708305358887f64]),
(device, burn::tensor::DType::F32),
),
device.clone(),
false,
[1].into(),
);
let linear26 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
Self {
layernormalization8,
conv1d4,
conv1d5,
conv1d6,
layernormalization9,
linear18,
linear19,
constant1348,
layernormalization10,
layernormalization11,
linear20,
linear21,
constant1349,
layernormalization12,
linear22,
linear23,
linear24,
linear25,
constant72,
constant73,
constant1395,
linear26,
device: device.clone(),
}
}
#[allow(clippy::let_and_return, clippy::approx_constant)]
pub fn forward(
&self,
add22_out1: Tensor<3>,
unsqueeze120_out1: Tensor<3, Bool>,
slice1_out1: Tensor<3>,
unsqueeze118_out1: Tensor<4, Bool>,
) -> Tensor<3> {
let layernormalization8_out1 = {
let dtype = add22_out1.clone().dtype();
self.layernormalization8
.forward(add22_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let transpose21_out1 = layernormalization8_out1.permute([0, 2, 1]);
let conv1d4_out1 = self.conv1d4.forward(transpose21_out1);
let split_tensors = conv1d4_out1.split_with_sizes([1024, 1024].into(), 1);
let [split2_out1, split2_out2] = split_tensors.try_into().unwrap();
let sigmoid6_out1 = burn::tensor::activation::sigmoid(split2_out2);
let mul55_out1 = split2_out1.mul(sigmoid6_out1);
let constant1339_out1 = 0f32;
let where24_out1 = mul55_out1.mask_fill(unsqueeze120_out1, constant1339_out1);
let pad4_out1 = where24_out1
.pad(
[(0usize, 0usize), (0usize, 0usize), (4usize, 4usize)],
burn::tensor::ops::PadMode::Constant(0f32),
);
let conv1d5_out1 = self.conv1d5.forward(pad4_out1);
let sigmoid7_out1 = burn::tensor::activation::sigmoid(conv1d5_out1.clone());
let mul56_out1 = conv1d5_out1.mul(sigmoid7_out1);
let conv1d6_out1 = self.conv1d6.forward(mul56_out1);
let transpose23_out1 = conv1d6_out1.permute([0, 2, 1]);
let add23_out1 = add22_out1.add(transpose23_out1);
let layernormalization9_out1 = {
let dtype = add23_out1.clone().dtype();
self.layernormalization9
.forward(add23_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let linear18_out1 = self.linear18.forward(layernormalization9_out1);
let sigmoid8_out1 = burn::tensor::activation::sigmoid(linear18_out1.clone());
let mul57_out1 = linear18_out1.mul(sigmoid8_out1);
let linear19_out1 = self.linear19.forward(mul57_out1);
let constant1348_out1 = self.constant1348.val();
let mul58_out1 = linear19_out1
.mul((constant1348_out1).unsqueeze_dims(&[0isize, 1isize]));
let add24_out1 = add23_out1.add(mul58_out1);
let layernormalization10_out1 = {
let dtype = add24_out1.dtype();
self.layernormalization10
.forward(add24_out1.cast(burn::tensor::DType::F32))
.cast(dtype)
};
let layernormalization11_out1 = {
let dtype = layernormalization10_out1.clone().dtype();
self.layernormalization11
.forward(
layernormalization10_out1.clone().cast(burn::tensor::DType::F32),
)
.cast(dtype)
};
let linear20_out1 = self.linear20.forward(layernormalization11_out1);
let sigmoid9_out1 = burn::tensor::activation::sigmoid(linear20_out1.clone());
let mul59_out1 = linear20_out1.mul(sigmoid9_out1);
let linear21_out1 = self.linear21.forward(mul59_out1);
let constant1349_out1 = self.constant1349.val();
let mul60_out1 = linear21_out1
.mul((constant1349_out1).unsqueeze_dims(&[0isize, 1isize]));
let add25_out1 = layernormalization10_out1.add(mul60_out1);
let layernormalization12_out1 = {
let dtype = add25_out1.clone().dtype();
self.layernormalization12
.forward(add25_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let shape76_out1: [i64; 3] = {
let axes = &layernormalization12_out1.clone().dims()[0..3];
let mut output = [0i64; 3];
for i in 0..3 {
output[i] = axes[i] as i64;
}
output
};
let gather55_out1 = shape76_out1[0] as i64;
let linear22_out1 = self.linear22.forward(layernormalization12_out1.clone());
let unsqueeze133_out1 = [gather55_out1 as i64];
let constant1354_out1: [i64; 1] = [128i64];
let constant1352_out1: [i64; 1] = [-1i64];
let constant1353_out1: [i64; 1] = [8i64];
let concat63_out1: [i64; 4usize] = [
&unsqueeze133_out1[..],
&constant1352_out1[..],
&constant1353_out1[..],
&constant1354_out1[..],
]
.concat()
.try_into()
.unwrap();
let unsqueeze134_out1 = [gather55_out1 as i64];
let constant1358_out1: [i64; 1] = [128i64];
let constant1356_out1: [i64; 1] = [-1i64];
let constant1357_out1: [i64; 1] = [8i64];
let concat64_out1: [i64; 4usize] = [
&unsqueeze134_out1[..],
&constant1356_out1[..],
&constant1357_out1[..],
&constant1358_out1[..],
]
.concat()
.try_into()
.unwrap();
let unsqueeze135_out1 = [gather55_out1 as i64];
let constant1362_out1: [i64; 1] = [128i64];
let constant1360_out1: [i64; 1] = [-1i64];
let constant1361_out1: [i64; 1] = [8i64];
let concat65_out1: [i64; 4usize] = [
&unsqueeze135_out1[..],
&constant1360_out1[..],
&constant1361_out1[..],
&constant1362_out1[..],
]
.concat()
.try_into()
.unwrap();
let reshape42_out1 = linear22_out1.reshape(concat63_out1);
let linear23_out1 = self.linear23.forward(layernormalization12_out1.clone());
let reshape43_out1 = linear23_out1.reshape(concat64_out1);
let linear24_out1 = self.linear24.forward(layernormalization12_out1);
let reshape44_out1 = linear24_out1.reshape(concat65_out1);
let transpose24_out1 = reshape44_out1.permute([0, 2, 1, 3]);
let linear25_out1 = self.linear25.forward(slice1_out1);
let reshape45_out1 = linear25_out1.reshape([1, -1, 8, 128]);
let constant72_out1 = self.constant72.val();
let add26_out1 = reshape42_out1
.clone()
.add((constant72_out1).unsqueeze_dims(&[0isize, 1isize]));
let transpose25_out1 = add26_out1.permute([0, 2, 1, 3]);
let constant73_out1 = self.constant73.val();
let add27_out1 = reshape42_out1
.add((constant73_out1).unsqueeze_dims(&[0isize, 1isize]));
let transpose26_out1 = add27_out1.permute([0, 2, 1, 3]);
let transpose27_out1 = reshape45_out1.permute([0, 2, 3, 1]);
let matmul32_out1 = transpose26_out1.matmul(transpose27_out1);
let shape77_out1: [i64; 4] = {
let axes = &matmul32_out1.clone().dims()[0..4];
let mut output = [0i64; 4];
for i in 0..4 {
output[i] = axes[i] as i64;
}
output
};
let gather56_out1 = shape77_out1[0] as i64;
let gather57_out1 = shape77_out1[1] as i64;
let gather58_out1 = shape77_out1[2] as i64;
let gather59_out1 = shape77_out1[3] as i64;
let pad5_out1 = matmul32_out1
.pad(
[(0usize, 0usize), (0usize, 0usize), (0usize, 0usize), (1usize, 0usize)],
burn::tensor::ops::PadMode::Constant(0f32),
);
let unsqueeze136_out1 = [gather56_out1 as i64];
let unsqueeze137_out1 = [gather57_out1 as i64];
let unsqueeze138_out1 = [gather58_out1 as i64];
let constant1377_out1: [i64; 1] = [-1i64];
let concat67_out1: [i64; 4usize] = [
&unsqueeze136_out1[..],
&unsqueeze137_out1[..],
&constant1377_out1[..],
&unsqueeze138_out1[..],
]
.concat()
.try_into()
.unwrap();
let reshape48_out1 = pad5_out1.reshape(concat67_out1);
let slice15_out1 = reshape48_out1.slice(s![.., .., 1.., ..]);
let unsqueeze139_out1 = [gather56_out1 as i64];
let unsqueeze140_out1 = [gather57_out1 as i64];
let unsqueeze141_out1 = [gather58_out1 as i64];
let unsqueeze142_out1 = [gather59_out1 as i64];
let concat68_out1: [i64; 4usize] = [
&unsqueeze139_out1[..],
&unsqueeze140_out1[..],
&unsqueeze141_out1[..],
&unsqueeze142_out1[..],
]
.concat()
.try_into()
.unwrap();
let reshape49_out1 = slice15_out1.reshape(concat68_out1);
let transpose29_out1 = reshape43_out1.permute([0, 2, 3, 1]);
let matmul33_out1 = transpose25_out1.matmul(transpose29_out1);
let shape81_out1: [i64; 4] = {
let axes = &matmul33_out1.clone().dims()[0..4];
let mut output = [0i64; 4];
for i in 0..4 {
output[i] = axes[i] as i64;
}
output
};
let slice16_out1: [i64; 1] = shape81_out1[3..4].try_into().unwrap();
let squeeze3_out1 = slice16_out1[0] as i64;
let unsqueeze143_out1 = [squeeze3_out1 as i64];
let slice17_out1 = reshape49_out1.slice(s![.., .., .., 0..unsqueeze143_out1[0]]);
let add28_out1 = matmul33_out1.add(slice17_out1);
let constant1395_out1 = self.constant1395.val();
let div6_out1 = add28_out1
.div((constant1395_out1).unsqueeze_dims(&[0isize, 1isize, 2isize]));
let shape82_out1: [i64; 4] = {
let axes = &transpose24_out1.clone().dims()[0..4];
let mut output = [0i64; 4];
for i in 0..4 {
output[i] = axes[i] as i64;
}
output
};
let gather60_out1 = shape82_out1[0] as i64;
let constant1397_out1 = -10000f32;
let where25_out1 = div6_out1
.mask_fill(unsqueeze118_out1.clone(), constant1397_out1);
let softmax3_out1 = burn::tensor::activation::softmax(where25_out1, 3);
let constant1398_out1 = 0f32;
let where26_out1 = softmax3_out1.mask_fill(unsqueeze118_out1, constant1398_out1);
let matmul34_out1 = where26_out1.matmul(transpose24_out1);
let transpose30_out1 = matmul34_out1.permute([0, 2, 1, 3]);
let unsqueeze144_out1 = [gather60_out1 as i64];
let constant1401_out1: [i64; 1] = [1024i64];
let constant1400_out1: [i64; 1] = [-1i64];
let concat69_out1: [i64; 3usize] = [
&unsqueeze144_out1[..],
&constant1400_out1[..],
&constant1401_out1[..],
]
.concat()
.try_into()
.unwrap();
let reshape50_out1 = transpose30_out1.reshape(concat69_out1);
let linear26_out1 = self.linear26.forward(reshape50_out1);
let add29_out1 = add25_out1.add(linear26_out1);
add29_out1
}
}
#[derive(Module, Debug)]
pub struct Submodule8 {
layernormalization13: LayerNorm,
conv1d7: Conv1d,
conv1d8: Conv1d,
conv1d9: Conv1d,
layernormalization14: LayerNorm,
linear27: Linear,
linear28: Linear,
constant1411: burn::module::Param<Tensor<1>>,
layernormalization15: LayerNorm,
layernormalization16: LayerNorm,
linear29: Linear,
linear30: Linear,
constant1412: burn::module::Param<Tensor<1>>,
layernormalization17: LayerNorm,
linear31: Linear,
linear32: Linear,
linear33: Linear,
linear34: Linear,
constant96: burn::module::Param<Tensor<2>>,
constant97: burn::module::Param<Tensor<2>>,
constant1458: burn::module::Param<Tensor<1>>,
linear35: Linear,
#[module(skip)]
device: Device,
}
impl Submodule8 {
#[allow(unused_variables)]
pub fn new(device: &Device) -> Self {
let layernormalization13 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let conv1d7 = Conv1dConfig::new(1024, 2048, 1)
.with_stride(1)
.with_padding(PaddingConfig1d::Valid)
.with_dilation(1)
.with_groups(1)
.with_bias(true)
.init(device);
let conv1d8 = Conv1dConfig::new(1024, 1024, 9)
.with_stride(1)
.with_padding(PaddingConfig1d::Valid)
.with_dilation(1)
.with_groups(1024)
.with_bias(true)
.init(device);
let conv1d9 = Conv1dConfig::new(1024, 1024, 1)
.with_stride(1)
.with_padding(PaddingConfig1d::Valid)
.with_dilation(1)
.with_groups(1)
.with_bias(true)
.init(device);
let layernormalization14 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let linear27 = LinearConfig::new(1024, 4096).with_bias(true).init(device);
let linear28 = LinearConfig::new(4096, 1024).with_bias(true).init(device);
let constant1411: 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([0.5f64]),
(device, burn::tensor::DType::F32),
),
device.clone(),
false,
[1].into(),
);
let layernormalization15 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let layernormalization16 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let linear29 = LinearConfig::new(1024, 4096).with_bias(true).init(device);
let linear30 = LinearConfig::new(4096, 1024).with_bias(true).init(device);
let constant1412: 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([0.5f64]),
(device, burn::tensor::DType::F32),
),
device.clone(),
false,
[1].into(),
);
let layernormalization17 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let linear31 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
let linear32 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
let linear33 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
let linear34 = LinearConfig::new(1024, 1024).with_bias(false).init(device);
let constant96: burn::module::Param<Tensor<2>> = burn::module::Param::uninitialized(
burn::module::ParamId::new(),
move |device, _require_grad| Tensor::<
2,
>::zeros([8, 128], (device, burn::tensor::DType::F32)),
device.clone(),
false,
[8, 128].into(),
);
let constant97: burn::module::Param<Tensor<2>> = burn::module::Param::uninitialized(
burn::module::ParamId::new(),
move |device, _require_grad| Tensor::<
2,
>::zeros([8, 128], (device, burn::tensor::DType::F32)),
device.clone(),
false,
[8, 128].into(),
);
let constant1458: 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([11.313708305358887f64]),
(device, burn::tensor::DType::F32),
),
device.clone(),
false,
[1].into(),
);
let linear35 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
Self {
layernormalization13,
conv1d7,
conv1d8,
conv1d9,
layernormalization14,
linear27,
linear28,
constant1411,
layernormalization15,
layernormalization16,
linear29,
linear30,
constant1412,
layernormalization17,
linear31,
linear32,
linear33,
linear34,
constant96,
constant97,
constant1458,
linear35,
device: device.clone(),
}
}
#[allow(clippy::let_and_return, clippy::approx_constant)]
pub fn forward(
&self,
add29_out1: Tensor<3>,
unsqueeze120_out1: Tensor<3, Bool>,
slice1_out1: Tensor<3>,
unsqueeze118_out1: Tensor<4, Bool>,
) -> Tensor<3> {
let layernormalization13_out1 = {
let dtype = add29_out1.clone().dtype();
self.layernormalization13
.forward(add29_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let transpose31_out1 = layernormalization13_out1.permute([0, 2, 1]);
let conv1d7_out1 = self.conv1d7.forward(transpose31_out1);
let split_tensors = conv1d7_out1.split_with_sizes([1024, 1024].into(), 1);
let [split3_out1, split3_out2] = split_tensors.try_into().unwrap();
let sigmoid10_out1 = burn::tensor::activation::sigmoid(split3_out2);
let mul61_out1 = split3_out1.mul(sigmoid10_out1);
let constant1402_out1 = 0f32;
let where27_out1 = mul61_out1.mask_fill(unsqueeze120_out1, constant1402_out1);
let pad6_out1 = where27_out1
.pad(
[(0usize, 0usize), (0usize, 0usize), (4usize, 4usize)],
burn::tensor::ops::PadMode::Constant(0f32),
);
let conv1d8_out1 = self.conv1d8.forward(pad6_out1);
let sigmoid11_out1 = burn::tensor::activation::sigmoid(conv1d8_out1.clone());
let mul62_out1 = conv1d8_out1.mul(sigmoid11_out1);
let conv1d9_out1 = self.conv1d9.forward(mul62_out1);
let transpose33_out1 = conv1d9_out1.permute([0, 2, 1]);
let add30_out1 = add29_out1.add(transpose33_out1);
let layernormalization14_out1 = {
let dtype = add30_out1.clone().dtype();
self.layernormalization14
.forward(add30_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let linear27_out1 = self.linear27.forward(layernormalization14_out1);
let sigmoid12_out1 = burn::tensor::activation::sigmoid(linear27_out1.clone());
let mul63_out1 = linear27_out1.mul(sigmoid12_out1);
let linear28_out1 = self.linear28.forward(mul63_out1);
let constant1411_out1 = self.constant1411.val();
let mul64_out1 = linear28_out1
.mul((constant1411_out1).unsqueeze_dims(&[0isize, 1isize]));
let add31_out1 = add30_out1.add(mul64_out1);
let layernormalization15_out1 = {
let dtype = add31_out1.dtype();
self.layernormalization15
.forward(add31_out1.cast(burn::tensor::DType::F32))
.cast(dtype)
};
let layernormalization16_out1 = {
let dtype = layernormalization15_out1.clone().dtype();
self.layernormalization16
.forward(
layernormalization15_out1.clone().cast(burn::tensor::DType::F32),
)
.cast(dtype)
};
let linear29_out1 = self.linear29.forward(layernormalization16_out1);
let sigmoid13_out1 = burn::tensor::activation::sigmoid(linear29_out1.clone());
let mul65_out1 = linear29_out1.mul(sigmoid13_out1);
let linear30_out1 = self.linear30.forward(mul65_out1);
let constant1412_out1 = self.constant1412.val();
let mul66_out1 = linear30_out1
.mul((constant1412_out1).unsqueeze_dims(&[0isize, 1isize]));
let add32_out1 = layernormalization15_out1.add(mul66_out1);
let layernormalization17_out1 = {
let dtype = add32_out1.clone().dtype();
self.layernormalization17
.forward(add32_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let shape83_out1: [i64; 3] = {
let axes = &layernormalization17_out1.clone().dims()[0..3];
let mut output = [0i64; 3];
for i in 0..3 {
output[i] = axes[i] as i64;
}
output
};
let gather61_out1 = shape83_out1[0] as i64;
let linear31_out1 = self.linear31.forward(layernormalization17_out1.clone());
let unsqueeze145_out1 = [gather61_out1 as i64];
let constant1417_out1: [i64; 1] = [128i64];
let constant1415_out1: [i64; 1] = [-1i64];
let constant1416_out1: [i64; 1] = [8i64];
let concat71_out1: [i64; 4usize] = [
&unsqueeze145_out1[..],
&constant1415_out1[..],
&constant1416_out1[..],
&constant1417_out1[..],
]
.concat()
.try_into()
.unwrap();
let unsqueeze146_out1 = [gather61_out1 as i64];
let constant1421_out1: [i64; 1] = [128i64];
let constant1419_out1: [i64; 1] = [-1i64];
let constant1420_out1: [i64; 1] = [8i64];
let concat72_out1: [i64; 4usize] = [
&unsqueeze146_out1[..],
&constant1419_out1[..],
&constant1420_out1[..],
&constant1421_out1[..],
]
.concat()
.try_into()
.unwrap();
let unsqueeze147_out1 = [gather61_out1 as i64];
let constant1425_out1: [i64; 1] = [128i64];
let constant1423_out1: [i64; 1] = [-1i64];
let constant1424_out1: [i64; 1] = [8i64];
let concat73_out1: [i64; 4usize] = [
&unsqueeze147_out1[..],
&constant1423_out1[..],
&constant1424_out1[..],
&constant1425_out1[..],
]
.concat()
.try_into()
.unwrap();
let reshape53_out1 = linear31_out1.reshape(concat71_out1);
let linear32_out1 = self.linear32.forward(layernormalization17_out1.clone());
let reshape54_out1 = linear32_out1.reshape(concat72_out1);
let linear33_out1 = self.linear33.forward(layernormalization17_out1);
let reshape55_out1 = linear33_out1.reshape(concat73_out1);
let transpose34_out1 = reshape55_out1.permute([0, 2, 1, 3]);
let linear34_out1 = self.linear34.forward(slice1_out1);
let reshape56_out1 = linear34_out1.reshape([1, -1, 8, 128]);
let constant96_out1 = self.constant96.val();
let add33_out1 = reshape53_out1
.clone()
.add((constant96_out1).unsqueeze_dims(&[0isize, 1isize]));
let transpose35_out1 = add33_out1.permute([0, 2, 1, 3]);
let constant97_out1 = self.constant97.val();
let add34_out1 = reshape53_out1
.add((constant97_out1).unsqueeze_dims(&[0isize, 1isize]));
let transpose36_out1 = add34_out1.permute([0, 2, 1, 3]);
let transpose37_out1 = reshape56_out1.permute([0, 2, 3, 1]);
let matmul44_out1 = transpose36_out1.matmul(transpose37_out1);
let shape84_out1: [i64; 4] = {
let axes = &matmul44_out1.clone().dims()[0..4];
let mut output = [0i64; 4];
for i in 0..4 {
output[i] = axes[i] as i64;
}
output
};
let gather62_out1 = shape84_out1[0] as i64;
let gather63_out1 = shape84_out1[1] as i64;
let gather64_out1 = shape84_out1[2] as i64;
let gather65_out1 = shape84_out1[3] as i64;
let pad7_out1 = matmul44_out1
.pad(
[(0usize, 0usize), (0usize, 0usize), (0usize, 0usize), (1usize, 0usize)],
burn::tensor::ops::PadMode::Constant(0f32),
);
let unsqueeze148_out1 = [gather62_out1 as i64];
let unsqueeze149_out1 = [gather63_out1 as i64];
let unsqueeze150_out1 = [gather64_out1 as i64];
let constant1440_out1: [i64; 1] = [-1i64];
let concat75_out1: [i64; 4usize] = [
&unsqueeze148_out1[..],
&unsqueeze149_out1[..],
&constant1440_out1[..],
&unsqueeze150_out1[..],
]
.concat()
.try_into()
.unwrap();
let reshape59_out1 = pad7_out1.reshape(concat75_out1);
let slice20_out1 = reshape59_out1.slice(s![.., .., 1.., ..]);
let unsqueeze151_out1 = [gather62_out1 as i64];
let unsqueeze152_out1 = [gather63_out1 as i64];
let unsqueeze153_out1 = [gather64_out1 as i64];
let unsqueeze154_out1 = [gather65_out1 as i64];
let concat76_out1: [i64; 4usize] = [
&unsqueeze151_out1[..],
&unsqueeze152_out1[..],
&unsqueeze153_out1[..],
&unsqueeze154_out1[..],
]
.concat()
.try_into()
.unwrap();
let reshape60_out1 = slice20_out1.reshape(concat76_out1);
let transpose39_out1 = reshape54_out1.permute([0, 2, 3, 1]);
let matmul45_out1 = transpose35_out1.matmul(transpose39_out1);
let shape88_out1: [i64; 4] = {
let axes = &matmul45_out1.clone().dims()[0..4];
let mut output = [0i64; 4];
for i in 0..4 {
output[i] = axes[i] as i64;
}
output
};
let slice21_out1: [i64; 1] = shape88_out1[3..4].try_into().unwrap();
let squeeze4_out1 = slice21_out1[0] as i64;
let unsqueeze155_out1 = [squeeze4_out1 as i64];
let slice22_out1 = reshape60_out1.slice(s![.., .., .., 0..unsqueeze155_out1[0]]);
let add35_out1 = matmul45_out1.add(slice22_out1);
let constant1458_out1 = self.constant1458.val();
let div7_out1 = add35_out1
.div((constant1458_out1).unsqueeze_dims(&[0isize, 1isize, 2isize]));
let shape89_out1: [i64; 4] = {
let axes = &transpose34_out1.clone().dims()[0..4];
let mut output = [0i64; 4];
for i in 0..4 {
output[i] = axes[i] as i64;
}
output
};
let gather66_out1 = shape89_out1[0] as i64;
let constant1460_out1 = -10000f32;
let where28_out1 = div7_out1
.mask_fill(unsqueeze118_out1.clone(), constant1460_out1);
let softmax4_out1 = burn::tensor::activation::softmax(where28_out1, 3);
let constant1461_out1 = 0f32;
let where29_out1 = softmax4_out1.mask_fill(unsqueeze118_out1, constant1461_out1);
let matmul46_out1 = where29_out1.matmul(transpose34_out1);
let transpose40_out1 = matmul46_out1.permute([0, 2, 1, 3]);
let unsqueeze156_out1 = [gather66_out1 as i64];
let constant1464_out1: [i64; 1] = [1024i64];
let constant1463_out1: [i64; 1] = [-1i64];
let concat77_out1: [i64; 3usize] = [
&unsqueeze156_out1[..],
&constant1463_out1[..],
&constant1464_out1[..],
]
.concat()
.try_into()
.unwrap();
let reshape61_out1 = transpose40_out1.reshape(concat77_out1);
let linear35_out1 = self.linear35.forward(reshape61_out1);
let add36_out1 = add32_out1.add(linear35_out1);
add36_out1
}
}
#[derive(Module, Debug)]
pub struct Submodule9 {
layernormalization18: LayerNorm,
conv1d10: Conv1d,
conv1d11: Conv1d,
conv1d12: Conv1d,
layernormalization19: LayerNorm,
linear36: Linear,
linear37: Linear,
constant1474: burn::module::Param<Tensor<1>>,
layernormalization20: LayerNorm,
layernormalization21: LayerNorm,
linear38: Linear,
linear39: Linear,
constant1475: burn::module::Param<Tensor<1>>,
layernormalization22: LayerNorm,
linear40: Linear,
linear41: Linear,
linear42: Linear,
linear43: Linear,
constant120: burn::module::Param<Tensor<2>>,
constant121: burn::module::Param<Tensor<2>>,
constant1521: burn::module::Param<Tensor<1>>,
linear44: Linear,
#[module(skip)]
device: Device,
}
impl Submodule9 {
#[allow(unused_variables)]
pub fn new(device: &Device) -> Self {
let layernormalization18 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let conv1d10 = Conv1dConfig::new(1024, 2048, 1)
.with_stride(1)
.with_padding(PaddingConfig1d::Valid)
.with_dilation(1)
.with_groups(1)
.with_bias(true)
.init(device);
let conv1d11 = Conv1dConfig::new(1024, 1024, 9)
.with_stride(1)
.with_padding(PaddingConfig1d::Valid)
.with_dilation(1)
.with_groups(1024)
.with_bias(true)
.init(device);
let conv1d12 = Conv1dConfig::new(1024, 1024, 1)
.with_stride(1)
.with_padding(PaddingConfig1d::Valid)
.with_dilation(1)
.with_groups(1)
.with_bias(true)
.init(device);
let layernormalization19 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let linear36 = LinearConfig::new(1024, 4096).with_bias(true).init(device);
let linear37 = LinearConfig::new(4096, 1024).with_bias(true).init(device);
let constant1474: 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([0.5f64]),
(device, burn::tensor::DType::F32),
),
device.clone(),
false,
[1].into(),
);
let layernormalization20 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let layernormalization21 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let linear38 = LinearConfig::new(1024, 4096).with_bias(true).init(device);
let linear39 = LinearConfig::new(4096, 1024).with_bias(true).init(device);
let constant1475: 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([0.5f64]),
(device, burn::tensor::DType::F32),
),
device.clone(),
false,
[1].into(),
);
let layernormalization22 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let linear40 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
let linear41 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
let linear42 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
let linear43 = LinearConfig::new(1024, 1024).with_bias(false).init(device);
let constant120: burn::module::Param<Tensor<2>> = burn::module::Param::uninitialized(
burn::module::ParamId::new(),
move |device, _require_grad| Tensor::<
2,
>::zeros([8, 128], (device, burn::tensor::DType::F32)),
device.clone(),
false,
[8, 128].into(),
);
let constant121: burn::module::Param<Tensor<2>> = burn::module::Param::uninitialized(
burn::module::ParamId::new(),
move |device, _require_grad| Tensor::<
2,
>::zeros([8, 128], (device, burn::tensor::DType::F32)),
device.clone(),
false,
[8, 128].into(),
);
let constant1521: 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([11.313708305358887f64]),
(device, burn::tensor::DType::F32),
),
device.clone(),
false,
[1].into(),
);
let linear44 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
Self {
layernormalization18,
conv1d10,
conv1d11,
conv1d12,
layernormalization19,
linear36,
linear37,
constant1474,
layernormalization20,
layernormalization21,
linear38,
linear39,
constant1475,
layernormalization22,
linear40,
linear41,
linear42,
linear43,
constant120,
constant121,
constant1521,
linear44,
device: device.clone(),
}
}
#[allow(clippy::let_and_return, clippy::approx_constant)]
pub fn forward(
&self,
add36_out1: Tensor<3>,
unsqueeze120_out1: Tensor<3, Bool>,
slice1_out1: Tensor<3>,
unsqueeze118_out1: Tensor<4, Bool>,
) -> Tensor<3> {
let layernormalization18_out1 = {
let dtype = add36_out1.clone().dtype();
self.layernormalization18
.forward(add36_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let transpose41_out1 = layernormalization18_out1.permute([0, 2, 1]);
let conv1d10_out1 = self.conv1d10.forward(transpose41_out1);
let split_tensors = conv1d10_out1.split_with_sizes([1024, 1024].into(), 1);
let [split4_out1, split4_out2] = split_tensors.try_into().unwrap();
let sigmoid14_out1 = burn::tensor::activation::sigmoid(split4_out2);
let mul67_out1 = split4_out1.mul(sigmoid14_out1);
let constant1465_out1 = 0f32;
let where30_out1 = mul67_out1.mask_fill(unsqueeze120_out1, constant1465_out1);
let pad8_out1 = where30_out1
.pad(
[(0usize, 0usize), (0usize, 0usize), (4usize, 4usize)],
burn::tensor::ops::PadMode::Constant(0f32),
);
let conv1d11_out1 = self.conv1d11.forward(pad8_out1);
let sigmoid15_out1 = burn::tensor::activation::sigmoid(conv1d11_out1.clone());
let mul68_out1 = conv1d11_out1.mul(sigmoid15_out1);
let conv1d12_out1 = self.conv1d12.forward(mul68_out1);
let transpose43_out1 = conv1d12_out1.permute([0, 2, 1]);
let add37_out1 = add36_out1.add(transpose43_out1);
let layernormalization19_out1 = {
let dtype = add37_out1.clone().dtype();
self.layernormalization19
.forward(add37_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let linear36_out1 = self.linear36.forward(layernormalization19_out1);
let sigmoid16_out1 = burn::tensor::activation::sigmoid(linear36_out1.clone());
let mul69_out1 = linear36_out1.mul(sigmoid16_out1);
let linear37_out1 = self.linear37.forward(mul69_out1);
let constant1474_out1 = self.constant1474.val();
let mul70_out1 = linear37_out1
.mul((constant1474_out1).unsqueeze_dims(&[0isize, 1isize]));
let add38_out1 = add37_out1.add(mul70_out1);
let layernormalization20_out1 = {
let dtype = add38_out1.dtype();
self.layernormalization20
.forward(add38_out1.cast(burn::tensor::DType::F32))
.cast(dtype)
};
let layernormalization21_out1 = {
let dtype = layernormalization20_out1.clone().dtype();
self.layernormalization21
.forward(
layernormalization20_out1.clone().cast(burn::tensor::DType::F32),
)
.cast(dtype)
};
let linear38_out1 = self.linear38.forward(layernormalization21_out1);
let sigmoid17_out1 = burn::tensor::activation::sigmoid(linear38_out1.clone());
let mul71_out1 = linear38_out1.mul(sigmoid17_out1);
let linear39_out1 = self.linear39.forward(mul71_out1);
let constant1475_out1 = self.constant1475.val();
let mul72_out1 = linear39_out1
.mul((constant1475_out1).unsqueeze_dims(&[0isize, 1isize]));
let add39_out1 = layernormalization20_out1.add(mul72_out1);
let layernormalization22_out1 = {
let dtype = add39_out1.clone().dtype();
self.layernormalization22
.forward(add39_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let shape90_out1: [i64; 3] = {
let axes = &layernormalization22_out1.clone().dims()[0..3];
let mut output = [0i64; 3];
for i in 0..3 {
output[i] = axes[i] as i64;
}
output
};
let gather67_out1 = shape90_out1[0] as i64;
let linear40_out1 = self.linear40.forward(layernormalization22_out1.clone());
let unsqueeze157_out1 = [gather67_out1 as i64];
let constant1480_out1: [i64; 1] = [128i64];
let constant1478_out1: [i64; 1] = [-1i64];
let constant1479_out1: [i64; 1] = [8i64];
let concat79_out1: [i64; 4usize] = [
&unsqueeze157_out1[..],
&constant1478_out1[..],
&constant1479_out1[..],
&constant1480_out1[..],
]
.concat()
.try_into()
.unwrap();
let unsqueeze158_out1 = [gather67_out1 as i64];
let constant1484_out1: [i64; 1] = [128i64];
let constant1482_out1: [i64; 1] = [-1i64];
let constant1483_out1: [i64; 1] = [8i64];
let concat80_out1: [i64; 4usize] = [
&unsqueeze158_out1[..],
&constant1482_out1[..],
&constant1483_out1[..],
&constant1484_out1[..],
]
.concat()
.try_into()
.unwrap();
let unsqueeze159_out1 = [gather67_out1 as i64];
let constant1488_out1: [i64; 1] = [128i64];
let constant1486_out1: [i64; 1] = [-1i64];
let constant1487_out1: [i64; 1] = [8i64];
let concat81_out1: [i64; 4usize] = [
&unsqueeze159_out1[..],
&constant1486_out1[..],
&constant1487_out1[..],
&constant1488_out1[..],
]
.concat()
.try_into()
.unwrap();
let reshape64_out1 = linear40_out1.reshape(concat79_out1);
let linear41_out1 = self.linear41.forward(layernormalization22_out1.clone());
let reshape65_out1 = linear41_out1.reshape(concat80_out1);
let linear42_out1 = self.linear42.forward(layernormalization22_out1);
let reshape66_out1 = linear42_out1.reshape(concat81_out1);
let transpose44_out1 = reshape66_out1.permute([0, 2, 1, 3]);
let linear43_out1 = self.linear43.forward(slice1_out1);
let reshape67_out1 = linear43_out1.reshape([1, -1, 8, 128]);
let constant120_out1 = self.constant120.val();
let add40_out1 = reshape64_out1
.clone()
.add((constant120_out1).unsqueeze_dims(&[0isize, 1isize]));
let transpose45_out1 = add40_out1.permute([0, 2, 1, 3]);
let constant121_out1 = self.constant121.val();
let add41_out1 = reshape64_out1
.add((constant121_out1).unsqueeze_dims(&[0isize, 1isize]));
let transpose46_out1 = add41_out1.permute([0, 2, 1, 3]);
let transpose47_out1 = reshape67_out1.permute([0, 2, 3, 1]);
let matmul56_out1 = transpose46_out1.matmul(transpose47_out1);
let shape91_out1: [i64; 4] = {
let axes = &matmul56_out1.clone().dims()[0..4];
let mut output = [0i64; 4];
for i in 0..4 {
output[i] = axes[i] as i64;
}
output
};
let gather68_out1 = shape91_out1[0] as i64;
let gather69_out1 = shape91_out1[1] as i64;
let gather70_out1 = shape91_out1[2] as i64;
let gather71_out1 = shape91_out1[3] as i64;
let pad9_out1 = matmul56_out1
.pad(
[(0usize, 0usize), (0usize, 0usize), (0usize, 0usize), (1usize, 0usize)],
burn::tensor::ops::PadMode::Constant(0f32),
);
let unsqueeze160_out1 = [gather68_out1 as i64];
let unsqueeze161_out1 = [gather69_out1 as i64];
let unsqueeze162_out1 = [gather70_out1 as i64];
let constant1503_out1: [i64; 1] = [-1i64];
let concat83_out1: [i64; 4usize] = [
&unsqueeze160_out1[..],
&unsqueeze161_out1[..],
&constant1503_out1[..],
&unsqueeze162_out1[..],
]
.concat()
.try_into()
.unwrap();
let reshape70_out1 = pad9_out1.reshape(concat83_out1);
let slice25_out1 = reshape70_out1.slice(s![.., .., 1.., ..]);
let unsqueeze163_out1 = [gather68_out1 as i64];
let unsqueeze164_out1 = [gather69_out1 as i64];
let unsqueeze165_out1 = [gather70_out1 as i64];
let unsqueeze166_out1 = [gather71_out1 as i64];
let concat84_out1: [i64; 4usize] = [
&unsqueeze163_out1[..],
&unsqueeze164_out1[..],
&unsqueeze165_out1[..],
&unsqueeze166_out1[..],
]
.concat()
.try_into()
.unwrap();
let reshape71_out1 = slice25_out1.reshape(concat84_out1);
let transpose49_out1 = reshape65_out1.permute([0, 2, 3, 1]);
let matmul57_out1 = transpose45_out1.matmul(transpose49_out1);
let shape95_out1: [i64; 4] = {
let axes = &matmul57_out1.clone().dims()[0..4];
let mut output = [0i64; 4];
for i in 0..4 {
output[i] = axes[i] as i64;
}
output
};
let slice26_out1: [i64; 1] = shape95_out1[3..4].try_into().unwrap();
let squeeze5_out1 = slice26_out1[0] as i64;
let unsqueeze167_out1 = [squeeze5_out1 as i64];
let slice27_out1 = reshape71_out1.slice(s![.., .., .., 0..unsqueeze167_out1[0]]);
let add42_out1 = matmul57_out1.add(slice27_out1);
let constant1521_out1 = self.constant1521.val();
let div8_out1 = add42_out1
.div((constant1521_out1).unsqueeze_dims(&[0isize, 1isize, 2isize]));
let shape96_out1: [i64; 4] = {
let axes = &transpose44_out1.clone().dims()[0..4];
let mut output = [0i64; 4];
for i in 0..4 {
output[i] = axes[i] as i64;
}
output
};
let gather72_out1 = shape96_out1[0] as i64;
let constant1523_out1 = -10000f32;
let where31_out1 = div8_out1
.mask_fill(unsqueeze118_out1.clone(), constant1523_out1);
let softmax5_out1 = burn::tensor::activation::softmax(where31_out1, 3);
let constant1524_out1 = 0f32;
let where32_out1 = softmax5_out1.mask_fill(unsqueeze118_out1, constant1524_out1);
let matmul58_out1 = where32_out1.matmul(transpose44_out1);
let transpose50_out1 = matmul58_out1.permute([0, 2, 1, 3]);
let unsqueeze168_out1 = [gather72_out1 as i64];
let constant1527_out1: [i64; 1] = [1024i64];
let constant1526_out1: [i64; 1] = [-1i64];
let concat85_out1: [i64; 3usize] = [
&unsqueeze168_out1[..],
&constant1526_out1[..],
&constant1527_out1[..],
]
.concat()
.try_into()
.unwrap();
let reshape72_out1 = transpose50_out1.reshape(concat85_out1);
let linear44_out1 = self.linear44.forward(reshape72_out1);
let add43_out1 = add39_out1.add(linear44_out1);
add43_out1
}
}
#[derive(Module, Debug)]
pub struct Submodule10 {
layernormalization23: LayerNorm,
conv1d13: Conv1d,
conv1d14: Conv1d,
conv1d15: Conv1d,
layernormalization24: LayerNorm,
linear45: Linear,
linear46: Linear,
constant1537: burn::module::Param<Tensor<1>>,
layernormalization25: LayerNorm,
layernormalization26: LayerNorm,
linear47: Linear,
linear48: Linear,
constant1538: burn::module::Param<Tensor<1>>,
layernormalization27: LayerNorm,
linear49: Linear,
linear50: Linear,
linear51: Linear,
linear52: Linear,
constant144: burn::module::Param<Tensor<2>>,
constant145: burn::module::Param<Tensor<2>>,
constant1584: burn::module::Param<Tensor<1>>,
linear53: Linear,
#[module(skip)]
device: Device,
}
impl Submodule10 {
#[allow(unused_variables)]
pub fn new(device: &Device) -> Self {
let layernormalization23 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let conv1d13 = Conv1dConfig::new(1024, 2048, 1)
.with_stride(1)
.with_padding(PaddingConfig1d::Valid)
.with_dilation(1)
.with_groups(1)
.with_bias(true)
.init(device);
let conv1d14 = Conv1dConfig::new(1024, 1024, 9)
.with_stride(1)
.with_padding(PaddingConfig1d::Valid)
.with_dilation(1)
.with_groups(1024)
.with_bias(true)
.init(device);
let conv1d15 = Conv1dConfig::new(1024, 1024, 1)
.with_stride(1)
.with_padding(PaddingConfig1d::Valid)
.with_dilation(1)
.with_groups(1)
.with_bias(true)
.init(device);
let layernormalization24 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let linear45 = LinearConfig::new(1024, 4096).with_bias(true).init(device);
let linear46 = LinearConfig::new(4096, 1024).with_bias(true).init(device);
let constant1537: 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([0.5f64]),
(device, burn::tensor::DType::F32),
),
device.clone(),
false,
[1].into(),
);
let layernormalization25 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let layernormalization26 = 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 constant1538: 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([0.5f64]),
(device, burn::tensor::DType::F32),
),
device.clone(),
false,
[1].into(),
);
let layernormalization27 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let linear49 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
let linear50 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
let linear51 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
let linear52 = LinearConfig::new(1024, 1024).with_bias(false).init(device);
let constant144: burn::module::Param<Tensor<2>> = burn::module::Param::uninitialized(
burn::module::ParamId::new(),
move |device, _require_grad| Tensor::<
2,
>::zeros([8, 128], (device, burn::tensor::DType::F32)),
device.clone(),
false,
[8, 128].into(),
);
let constant145: burn::module::Param<Tensor<2>> = burn::module::Param::uninitialized(
burn::module::ParamId::new(),
move |device, _require_grad| Tensor::<
2,
>::zeros([8, 128], (device, burn::tensor::DType::F32)),
device.clone(),
false,
[8, 128].into(),
);
let constant1584: 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([11.313708305358887f64]),
(device, burn::tensor::DType::F32),
),
device.clone(),
false,
[1].into(),
);
let linear53 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
Self {
layernormalization23,
conv1d13,
conv1d14,
conv1d15,
layernormalization24,
linear45,
linear46,
constant1537,
layernormalization25,
layernormalization26,
linear47,
linear48,
constant1538,
layernormalization27,
linear49,
linear50,
linear51,
linear52,
constant144,
constant145,
constant1584,
linear53,
device: device.clone(),
}
}
#[allow(clippy::let_and_return, clippy::approx_constant)]
pub fn forward(
&self,
add43_out1: Tensor<3>,
unsqueeze120_out1: Tensor<3, Bool>,
slice1_out1: Tensor<3>,
unsqueeze118_out1: Tensor<4, Bool>,
) -> Tensor<3> {
let layernormalization23_out1 = {
let dtype = add43_out1.clone().dtype();
self.layernormalization23
.forward(add43_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let transpose51_out1 = layernormalization23_out1.permute([0, 2, 1]);
let conv1d13_out1 = self.conv1d13.forward(transpose51_out1);
let split_tensors = conv1d13_out1.split_with_sizes([1024, 1024].into(), 1);
let [split5_out1, split5_out2] = split_tensors.try_into().unwrap();
let sigmoid18_out1 = burn::tensor::activation::sigmoid(split5_out2);
let mul73_out1 = split5_out1.mul(sigmoid18_out1);
let constant1528_out1 = 0f32;
let where33_out1 = mul73_out1.mask_fill(unsqueeze120_out1, constant1528_out1);
let pad10_out1 = where33_out1
.pad(
[(0usize, 0usize), (0usize, 0usize), (4usize, 4usize)],
burn::tensor::ops::PadMode::Constant(0f32),
);
let conv1d14_out1 = self.conv1d14.forward(pad10_out1);
let sigmoid19_out1 = burn::tensor::activation::sigmoid(conv1d14_out1.clone());
let mul74_out1 = conv1d14_out1.mul(sigmoid19_out1);
let conv1d15_out1 = self.conv1d15.forward(mul74_out1);
let transpose53_out1 = conv1d15_out1.permute([0, 2, 1]);
let add44_out1 = add43_out1.add(transpose53_out1);
let layernormalization24_out1 = {
let dtype = add44_out1.clone().dtype();
self.layernormalization24
.forward(add44_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let linear45_out1 = self.linear45.forward(layernormalization24_out1);
let sigmoid20_out1 = burn::tensor::activation::sigmoid(linear45_out1.clone());
let mul75_out1 = linear45_out1.mul(sigmoid20_out1);
let linear46_out1 = self.linear46.forward(mul75_out1);
let constant1537_out1 = self.constant1537.val();
let mul76_out1 = linear46_out1
.mul((constant1537_out1).unsqueeze_dims(&[0isize, 1isize]));
let add45_out1 = add44_out1.add(mul76_out1);
let layernormalization25_out1 = {
let dtype = add45_out1.dtype();
self.layernormalization25
.forward(add45_out1.cast(burn::tensor::DType::F32))
.cast(dtype)
};
let layernormalization26_out1 = {
let dtype = layernormalization25_out1.clone().dtype();
self.layernormalization26
.forward(
layernormalization25_out1.clone().cast(burn::tensor::DType::F32),
)
.cast(dtype)
};
let linear47_out1 = self.linear47.forward(layernormalization26_out1);
let sigmoid21_out1 = burn::tensor::activation::sigmoid(linear47_out1.clone());
let mul77_out1 = linear47_out1.mul(sigmoid21_out1);
let linear48_out1 = self.linear48.forward(mul77_out1);
let constant1538_out1 = self.constant1538.val();
let mul78_out1 = linear48_out1
.mul((constant1538_out1).unsqueeze_dims(&[0isize, 1isize]));
let add46_out1 = layernormalization25_out1.add(mul78_out1);
let layernormalization27_out1 = {
let dtype = add46_out1.clone().dtype();
self.layernormalization27
.forward(add46_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let shape97_out1: [i64; 3] = {
let axes = &layernormalization27_out1.clone().dims()[0..3];
let mut output = [0i64; 3];
for i in 0..3 {
output[i] = axes[i] as i64;
}
output
};
let gather73_out1 = shape97_out1[0] as i64;
let linear49_out1 = self.linear49.forward(layernormalization27_out1.clone());
let unsqueeze169_out1 = [gather73_out1 as i64];
let constant1543_out1: [i64; 1] = [128i64];
let constant1541_out1: [i64; 1] = [-1i64];
let constant1542_out1: [i64; 1] = [8i64];
let concat87_out1: [i64; 4usize] = [
&unsqueeze169_out1[..],
&constant1541_out1[..],
&constant1542_out1[..],
&constant1543_out1[..],
]
.concat()
.try_into()
.unwrap();
let unsqueeze170_out1 = [gather73_out1 as i64];
let constant1547_out1: [i64; 1] = [128i64];
let constant1545_out1: [i64; 1] = [-1i64];
let constant1546_out1: [i64; 1] = [8i64];
let concat88_out1: [i64; 4usize] = [
&unsqueeze170_out1[..],
&constant1545_out1[..],
&constant1546_out1[..],
&constant1547_out1[..],
]
.concat()
.try_into()
.unwrap();
let unsqueeze171_out1 = [gather73_out1 as i64];
let constant1551_out1: [i64; 1] = [128i64];
let constant1549_out1: [i64; 1] = [-1i64];
let constant1550_out1: [i64; 1] = [8i64];
let concat89_out1: [i64; 4usize] = [
&unsqueeze171_out1[..],
&constant1549_out1[..],
&constant1550_out1[..],
&constant1551_out1[..],
]
.concat()
.try_into()
.unwrap();
let reshape75_out1 = linear49_out1.reshape(concat87_out1);
let linear50_out1 = self.linear50.forward(layernormalization27_out1.clone());
let reshape76_out1 = linear50_out1.reshape(concat88_out1);
let linear51_out1 = self.linear51.forward(layernormalization27_out1);
let reshape77_out1 = linear51_out1.reshape(concat89_out1);
let transpose54_out1 = reshape77_out1.permute([0, 2, 1, 3]);
let linear52_out1 = self.linear52.forward(slice1_out1);
let reshape78_out1 = linear52_out1.reshape([1, -1, 8, 128]);
let constant144_out1 = self.constant144.val();
let add47_out1 = reshape75_out1
.clone()
.add((constant144_out1).unsqueeze_dims(&[0isize, 1isize]));
let transpose55_out1 = add47_out1.permute([0, 2, 1, 3]);
let constant145_out1 = self.constant145.val();
let add48_out1 = reshape75_out1
.add((constant145_out1).unsqueeze_dims(&[0isize, 1isize]));
let transpose56_out1 = add48_out1.permute([0, 2, 1, 3]);
let transpose57_out1 = reshape78_out1.permute([0, 2, 3, 1]);
let matmul68_out1 = transpose56_out1.matmul(transpose57_out1);
let shape98_out1: [i64; 4] = {
let axes = &matmul68_out1.clone().dims()[0..4];
let mut output = [0i64; 4];
for i in 0..4 {
output[i] = axes[i] as i64;
}
output
};
let gather74_out1 = shape98_out1[0] as i64;
let gather75_out1 = shape98_out1[1] as i64;
let gather76_out1 = shape98_out1[2] as i64;
let gather77_out1 = shape98_out1[3] as i64;
let pad11_out1 = matmul68_out1
.pad(
[(0usize, 0usize), (0usize, 0usize), (0usize, 0usize), (1usize, 0usize)],
burn::tensor::ops::PadMode::Constant(0f32),
);
let unsqueeze172_out1 = [gather74_out1 as i64];
let unsqueeze173_out1 = [gather75_out1 as i64];
let unsqueeze174_out1 = [gather76_out1 as i64];
let constant1566_out1: [i64; 1] = [-1i64];
let concat91_out1: [i64; 4usize] = [
&unsqueeze172_out1[..],
&unsqueeze173_out1[..],
&constant1566_out1[..],
&unsqueeze174_out1[..],
]
.concat()
.try_into()
.unwrap();
let reshape81_out1 = pad11_out1.reshape(concat91_out1);
let slice30_out1 = reshape81_out1.slice(s![.., .., 1.., ..]);
let unsqueeze175_out1 = [gather74_out1 as i64];
let unsqueeze176_out1 = [gather75_out1 as i64];
let unsqueeze177_out1 = [gather76_out1 as i64];
let unsqueeze178_out1 = [gather77_out1 as i64];
let concat92_out1: [i64; 4usize] = [
&unsqueeze175_out1[..],
&unsqueeze176_out1[..],
&unsqueeze177_out1[..],
&unsqueeze178_out1[..],
]
.concat()
.try_into()
.unwrap();
let reshape82_out1 = slice30_out1.reshape(concat92_out1);
let transpose59_out1 = reshape76_out1.permute([0, 2, 3, 1]);
let matmul69_out1 = transpose55_out1.matmul(transpose59_out1);
let shape102_out1: [i64; 4] = {
let axes = &matmul69_out1.clone().dims()[0..4];
let mut output = [0i64; 4];
for i in 0..4 {
output[i] = axes[i] as i64;
}
output
};
let slice31_out1: [i64; 1] = shape102_out1[3..4].try_into().unwrap();
let squeeze6_out1 = slice31_out1[0] as i64;
let unsqueeze179_out1 = [squeeze6_out1 as i64];
let slice32_out1 = reshape82_out1.slice(s![.., .., .., 0..unsqueeze179_out1[0]]);
let add49_out1 = matmul69_out1.add(slice32_out1);
let constant1584_out1 = self.constant1584.val();
let div9_out1 = add49_out1
.div((constant1584_out1).unsqueeze_dims(&[0isize, 1isize, 2isize]));
let shape103_out1: [i64; 4] = {
let axes = &transpose54_out1.clone().dims()[0..4];
let mut output = [0i64; 4];
for i in 0..4 {
output[i] = axes[i] as i64;
}
output
};
let gather78_out1 = shape103_out1[0] as i64;
let constant1586_out1 = -10000f32;
let where34_out1 = div9_out1
.mask_fill(unsqueeze118_out1.clone(), constant1586_out1);
let softmax6_out1 = burn::tensor::activation::softmax(where34_out1, 3);
let constant1587_out1 = 0f32;
let where35_out1 = softmax6_out1.mask_fill(unsqueeze118_out1, constant1587_out1);
let matmul70_out1 = where35_out1.matmul(transpose54_out1);
let transpose60_out1 = matmul70_out1.permute([0, 2, 1, 3]);
let unsqueeze180_out1 = [gather78_out1 as i64];
let constant1590_out1: [i64; 1] = [1024i64];
let constant1589_out1: [i64; 1] = [-1i64];
let concat93_out1: [i64; 3usize] = [
&unsqueeze180_out1[..],
&constant1589_out1[..],
&constant1590_out1[..],
]
.concat()
.try_into()
.unwrap();
let reshape83_out1 = transpose60_out1.reshape(concat93_out1);
let linear53_out1 = self.linear53.forward(reshape83_out1);
let add50_out1 = add46_out1.add(linear53_out1);
add50_out1
}
}
#[derive(Module, Debug)]
pub struct Submodule11 {
layernormalization28: LayerNorm,
conv1d16: Conv1d,
conv1d17: Conv1d,
conv1d18: Conv1d,
layernormalization29: LayerNorm,
linear54: Linear,
linear55: Linear,
constant1600: burn::module::Param<Tensor<1>>,
layernormalization30: LayerNorm,
layernormalization31: LayerNorm,
linear56: Linear,
linear57: Linear,
constant1601: burn::module::Param<Tensor<1>>,
layernormalization32: LayerNorm,
linear58: Linear,
linear59: Linear,
linear60: Linear,
linear61: Linear,
constant168: burn::module::Param<Tensor<2>>,
constant169: burn::module::Param<Tensor<2>>,
constant1647: burn::module::Param<Tensor<1>>,
linear62: Linear,
#[module(skip)]
device: Device,
}
impl Submodule11 {
#[allow(unused_variables)]
pub fn new(device: &Device) -> Self {
let layernormalization28 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let conv1d16 = Conv1dConfig::new(1024, 2048, 1)
.with_stride(1)
.with_padding(PaddingConfig1d::Valid)
.with_dilation(1)
.with_groups(1)
.with_bias(true)
.init(device);
let conv1d17 = Conv1dConfig::new(1024, 1024, 9)
.with_stride(1)
.with_padding(PaddingConfig1d::Valid)
.with_dilation(1)
.with_groups(1024)
.with_bias(true)
.init(device);
let conv1d18 = Conv1dConfig::new(1024, 1024, 1)
.with_stride(1)
.with_padding(PaddingConfig1d::Valid)
.with_dilation(1)
.with_groups(1)
.with_bias(true)
.init(device);
let layernormalization29 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let linear54 = LinearConfig::new(1024, 4096).with_bias(true).init(device);
let linear55 = LinearConfig::new(4096, 1024).with_bias(true).init(device);
let constant1600: 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([0.5f64]),
(device, burn::tensor::DType::F32),
),
device.clone(),
false,
[1].into(),
);
let layernormalization30 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let layernormalization31 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let linear56 = LinearConfig::new(1024, 4096).with_bias(true).init(device);
let linear57 = LinearConfig::new(4096, 1024).with_bias(true).init(device);
let constant1601: 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([0.5f64]),
(device, burn::tensor::DType::F32),
),
device.clone(),
false,
[1].into(),
);
let layernormalization32 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let linear58 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
let linear59 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
let linear60 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
let linear61 = LinearConfig::new(1024, 1024).with_bias(false).init(device);
let constant168: burn::module::Param<Tensor<2>> = burn::module::Param::uninitialized(
burn::module::ParamId::new(),
move |device, _require_grad| Tensor::<
2,
>::zeros([8, 128], (device, burn::tensor::DType::F32)),
device.clone(),
false,
[8, 128].into(),
);
let constant169: burn::module::Param<Tensor<2>> = burn::module::Param::uninitialized(
burn::module::ParamId::new(),
move |device, _require_grad| Tensor::<
2,
>::zeros([8, 128], (device, burn::tensor::DType::F32)),
device.clone(),
false,
[8, 128].into(),
);
let constant1647: 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([11.313708305358887f64]),
(device, burn::tensor::DType::F32),
),
device.clone(),
false,
[1].into(),
);
let linear62 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
Self {
layernormalization28,
conv1d16,
conv1d17,
conv1d18,
layernormalization29,
linear54,
linear55,
constant1600,
layernormalization30,
layernormalization31,
linear56,
linear57,
constant1601,
layernormalization32,
linear58,
linear59,
linear60,
linear61,
constant168,
constant169,
constant1647,
linear62,
device: device.clone(),
}
}
#[allow(clippy::let_and_return, clippy::approx_constant)]
pub fn forward(
&self,
add50_out1: Tensor<3>,
unsqueeze120_out1: Tensor<3, Bool>,
slice1_out1: Tensor<3>,
unsqueeze118_out1: Tensor<4, Bool>,
) -> Tensor<3> {
let layernormalization28_out1 = {
let dtype = add50_out1.clone().dtype();
self.layernormalization28
.forward(add50_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let transpose61_out1 = layernormalization28_out1.permute([0, 2, 1]);
let conv1d16_out1 = self.conv1d16.forward(transpose61_out1);
let split_tensors = conv1d16_out1.split_with_sizes([1024, 1024].into(), 1);
let [split6_out1, split6_out2] = split_tensors.try_into().unwrap();
let sigmoid22_out1 = burn::tensor::activation::sigmoid(split6_out2);
let mul79_out1 = split6_out1.mul(sigmoid22_out1);
let constant1591_out1 = 0f32;
let where36_out1 = mul79_out1.mask_fill(unsqueeze120_out1, constant1591_out1);
let pad12_out1 = where36_out1
.pad(
[(0usize, 0usize), (0usize, 0usize), (4usize, 4usize)],
burn::tensor::ops::PadMode::Constant(0f32),
);
let conv1d17_out1 = self.conv1d17.forward(pad12_out1);
let sigmoid23_out1 = burn::tensor::activation::sigmoid(conv1d17_out1.clone());
let mul80_out1 = conv1d17_out1.mul(sigmoid23_out1);
let conv1d18_out1 = self.conv1d18.forward(mul80_out1);
let transpose63_out1 = conv1d18_out1.permute([0, 2, 1]);
let add51_out1 = add50_out1.add(transpose63_out1);
let layernormalization29_out1 = {
let dtype = add51_out1.clone().dtype();
self.layernormalization29
.forward(add51_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let linear54_out1 = self.linear54.forward(layernormalization29_out1);
let sigmoid24_out1 = burn::tensor::activation::sigmoid(linear54_out1.clone());
let mul81_out1 = linear54_out1.mul(sigmoid24_out1);
let linear55_out1 = self.linear55.forward(mul81_out1);
let constant1600_out1 = self.constant1600.val();
let mul82_out1 = linear55_out1
.mul((constant1600_out1).unsqueeze_dims(&[0isize, 1isize]));
let add52_out1 = add51_out1.add(mul82_out1);
let layernormalization30_out1 = {
let dtype = add52_out1.dtype();
self.layernormalization30
.forward(add52_out1.cast(burn::tensor::DType::F32))
.cast(dtype)
};
let layernormalization31_out1 = {
let dtype = layernormalization30_out1.clone().dtype();
self.layernormalization31
.forward(
layernormalization30_out1.clone().cast(burn::tensor::DType::F32),
)
.cast(dtype)
};
let linear56_out1 = self.linear56.forward(layernormalization31_out1);
let sigmoid25_out1 = burn::tensor::activation::sigmoid(linear56_out1.clone());
let mul83_out1 = linear56_out1.mul(sigmoid25_out1);
let linear57_out1 = self.linear57.forward(mul83_out1);
let constant1601_out1 = self.constant1601.val();
let mul84_out1 = linear57_out1
.mul((constant1601_out1).unsqueeze_dims(&[0isize, 1isize]));
let add53_out1 = layernormalization30_out1.add(mul84_out1);
let layernormalization32_out1 = {
let dtype = add53_out1.clone().dtype();
self.layernormalization32
.forward(add53_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let shape104_out1: [i64; 3] = {
let axes = &layernormalization32_out1.clone().dims()[0..3];
let mut output = [0i64; 3];
for i in 0..3 {
output[i] = axes[i] as i64;
}
output
};
let gather79_out1 = shape104_out1[0] as i64;
let linear58_out1 = self.linear58.forward(layernormalization32_out1.clone());
let unsqueeze181_out1 = [gather79_out1 as i64];
let constant1606_out1: [i64; 1] = [128i64];
let constant1604_out1: [i64; 1] = [-1i64];
let constant1605_out1: [i64; 1] = [8i64];
let concat95_out1: [i64; 4usize] = [
&unsqueeze181_out1[..],
&constant1604_out1[..],
&constant1605_out1[..],
&constant1606_out1[..],
]
.concat()
.try_into()
.unwrap();
let unsqueeze182_out1 = [gather79_out1 as i64];
let constant1610_out1: [i64; 1] = [128i64];
let constant1608_out1: [i64; 1] = [-1i64];
let constant1609_out1: [i64; 1] = [8i64];
let concat96_out1: [i64; 4usize] = [
&unsqueeze182_out1[..],
&constant1608_out1[..],
&constant1609_out1[..],
&constant1610_out1[..],
]
.concat()
.try_into()
.unwrap();
let unsqueeze183_out1 = [gather79_out1 as i64];
let constant1614_out1: [i64; 1] = [128i64];
let constant1612_out1: [i64; 1] = [-1i64];
let constant1613_out1: [i64; 1] = [8i64];
let concat97_out1: [i64; 4usize] = [
&unsqueeze183_out1[..],
&constant1612_out1[..],
&constant1613_out1[..],
&constant1614_out1[..],
]
.concat()
.try_into()
.unwrap();
let reshape86_out1 = linear58_out1.reshape(concat95_out1);
let linear59_out1 = self.linear59.forward(layernormalization32_out1.clone());
let reshape87_out1 = linear59_out1.reshape(concat96_out1);
let linear60_out1 = self.linear60.forward(layernormalization32_out1);
let reshape88_out1 = linear60_out1.reshape(concat97_out1);
let transpose64_out1 = reshape88_out1.permute([0, 2, 1, 3]);
let linear61_out1 = self.linear61.forward(slice1_out1);
let reshape89_out1 = linear61_out1.reshape([1, -1, 8, 128]);
let constant168_out1 = self.constant168.val();
let add54_out1 = reshape86_out1
.clone()
.add((constant168_out1).unsqueeze_dims(&[0isize, 1isize]));
let transpose65_out1 = add54_out1.permute([0, 2, 1, 3]);
let constant169_out1 = self.constant169.val();
let add55_out1 = reshape86_out1
.add((constant169_out1).unsqueeze_dims(&[0isize, 1isize]));
let transpose66_out1 = add55_out1.permute([0, 2, 1, 3]);
let transpose67_out1 = reshape89_out1.permute([0, 2, 3, 1]);
let matmul80_out1 = transpose66_out1.matmul(transpose67_out1);
let shape105_out1: [i64; 4] = {
let axes = &matmul80_out1.clone().dims()[0..4];
let mut output = [0i64; 4];
for i in 0..4 {
output[i] = axes[i] as i64;
}
output
};
let gather80_out1 = shape105_out1[0] as i64;
let gather81_out1 = shape105_out1[1] as i64;
let gather82_out1 = shape105_out1[2] as i64;
let gather83_out1 = shape105_out1[3] as i64;
let pad13_out1 = matmul80_out1
.pad(
[(0usize, 0usize), (0usize, 0usize), (0usize, 0usize), (1usize, 0usize)],
burn::tensor::ops::PadMode::Constant(0f32),
);
let unsqueeze184_out1 = [gather80_out1 as i64];
let unsqueeze185_out1 = [gather81_out1 as i64];
let unsqueeze186_out1 = [gather82_out1 as i64];
let constant1629_out1: [i64; 1] = [-1i64];
let concat99_out1: [i64; 4usize] = [
&unsqueeze184_out1[..],
&unsqueeze185_out1[..],
&constant1629_out1[..],
&unsqueeze186_out1[..],
]
.concat()
.try_into()
.unwrap();
let reshape92_out1 = pad13_out1.reshape(concat99_out1);
let slice35_out1 = reshape92_out1.slice(s![.., .., 1.., ..]);
let unsqueeze187_out1 = [gather80_out1 as i64];
let unsqueeze188_out1 = [gather81_out1 as i64];
let unsqueeze189_out1 = [gather82_out1 as i64];
let unsqueeze190_out1 = [gather83_out1 as i64];
let concat100_out1: [i64; 4usize] = [
&unsqueeze187_out1[..],
&unsqueeze188_out1[..],
&unsqueeze189_out1[..],
&unsqueeze190_out1[..],
]
.concat()
.try_into()
.unwrap();
let reshape93_out1 = slice35_out1.reshape(concat100_out1);
let transpose69_out1 = reshape87_out1.permute([0, 2, 3, 1]);
let matmul81_out1 = transpose65_out1.matmul(transpose69_out1);
let shape109_out1: [i64; 4] = {
let axes = &matmul81_out1.clone().dims()[0..4];
let mut output = [0i64; 4];
for i in 0..4 {
output[i] = axes[i] as i64;
}
output
};
let slice36_out1: [i64; 1] = shape109_out1[3..4].try_into().unwrap();
let squeeze7_out1 = slice36_out1[0] as i64;
let unsqueeze191_out1 = [squeeze7_out1 as i64];
let slice37_out1 = reshape93_out1.slice(s![.., .., .., 0..unsqueeze191_out1[0]]);
let add56_out1 = matmul81_out1.add(slice37_out1);
let constant1647_out1 = self.constant1647.val();
let div10_out1 = add56_out1
.div((constant1647_out1).unsqueeze_dims(&[0isize, 1isize, 2isize]));
let shape110_out1: [i64; 4] = {
let axes = &transpose64_out1.clone().dims()[0..4];
let mut output = [0i64; 4];
for i in 0..4 {
output[i] = axes[i] as i64;
}
output
};
let gather84_out1 = shape110_out1[0] as i64;
let constant1649_out1 = -10000f32;
let where37_out1 = div10_out1
.mask_fill(unsqueeze118_out1.clone(), constant1649_out1);
let softmax7_out1 = burn::tensor::activation::softmax(where37_out1, 3);
let constant1650_out1 = 0f32;
let where38_out1 = softmax7_out1.mask_fill(unsqueeze118_out1, constant1650_out1);
let matmul82_out1 = where38_out1.matmul(transpose64_out1);
let transpose70_out1 = matmul82_out1.permute([0, 2, 1, 3]);
let unsqueeze192_out1 = [gather84_out1 as i64];
let constant1653_out1: [i64; 1] = [1024i64];
let constant1652_out1: [i64; 1] = [-1i64];
let concat101_out1: [i64; 3usize] = [
&unsqueeze192_out1[..],
&constant1652_out1[..],
&constant1653_out1[..],
]
.concat()
.try_into()
.unwrap();
let reshape94_out1 = transpose70_out1.reshape(concat101_out1);
let linear62_out1 = self.linear62.forward(reshape94_out1);
let add57_out1 = add53_out1.add(linear62_out1);
add57_out1
}
}
#[derive(Module, Debug)]
pub struct Submodule12 {
layernormalization33: LayerNorm,
conv1d19: Conv1d,
conv1d20: Conv1d,
conv1d21: Conv1d,
layernormalization34: LayerNorm,
linear63: Linear,
linear64: Linear,
constant1663: burn::module::Param<Tensor<1>>,
layernormalization35: LayerNorm,
layernormalization36: LayerNorm,
linear65: Linear,
linear66: Linear,
constant1664: burn::module::Param<Tensor<1>>,
layernormalization37: LayerNorm,
linear67: Linear,
linear68: Linear,
linear69: Linear,
linear70: Linear,
constant192: burn::module::Param<Tensor<2>>,
constant193: burn::module::Param<Tensor<2>>,
constant1710: burn::module::Param<Tensor<1>>,
linear71: Linear,
#[module(skip)]
device: Device,
}
impl Submodule12 {
#[allow(unused_variables)]
pub fn new(device: &Device) -> Self {
let layernormalization33 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let conv1d19 = Conv1dConfig::new(1024, 2048, 1)
.with_stride(1)
.with_padding(PaddingConfig1d::Valid)
.with_dilation(1)
.with_groups(1)
.with_bias(true)
.init(device);
let conv1d20 = Conv1dConfig::new(1024, 1024, 9)
.with_stride(1)
.with_padding(PaddingConfig1d::Valid)
.with_dilation(1)
.with_groups(1024)
.with_bias(true)
.init(device);
let conv1d21 = Conv1dConfig::new(1024, 1024, 1)
.with_stride(1)
.with_padding(PaddingConfig1d::Valid)
.with_dilation(1)
.with_groups(1)
.with_bias(true)
.init(device);
let layernormalization34 = 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 constant1663: 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([0.5f64]),
(device, burn::tensor::DType::F32),
),
device.clone(),
false,
[1].into(),
);
let layernormalization35 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let layernormalization36 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let linear65 = LinearConfig::new(1024, 4096).with_bias(true).init(device);
let linear66 = LinearConfig::new(4096, 1024).with_bias(true).init(device);
let constant1664: 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([0.5f64]),
(device, burn::tensor::DType::F32),
),
device.clone(),
false,
[1].into(),
);
let layernormalization37 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let linear67 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
let linear68 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
let linear69 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
let linear70 = LinearConfig::new(1024, 1024).with_bias(false).init(device);
let constant192: burn::module::Param<Tensor<2>> = burn::module::Param::uninitialized(
burn::module::ParamId::new(),
move |device, _require_grad| Tensor::<
2,
>::zeros([8, 128], (device, burn::tensor::DType::F32)),
device.clone(),
false,
[8, 128].into(),
);
let constant193: burn::module::Param<Tensor<2>> = burn::module::Param::uninitialized(
burn::module::ParamId::new(),
move |device, _require_grad| Tensor::<
2,
>::zeros([8, 128], (device, burn::tensor::DType::F32)),
device.clone(),
false,
[8, 128].into(),
);
let constant1710: 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([11.313708305358887f64]),
(device, burn::tensor::DType::F32),
),
device.clone(),
false,
[1].into(),
);
let linear71 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
Self {
layernormalization33,
conv1d19,
conv1d20,
conv1d21,
layernormalization34,
linear63,
linear64,
constant1663,
layernormalization35,
layernormalization36,
linear65,
linear66,
constant1664,
layernormalization37,
linear67,
linear68,
linear69,
linear70,
constant192,
constant193,
constant1710,
linear71,
device: device.clone(),
}
}
#[allow(clippy::let_and_return, clippy::approx_constant)]
pub fn forward(
&self,
add57_out1: Tensor<3>,
unsqueeze120_out1: Tensor<3, Bool>,
slice1_out1: Tensor<3>,
unsqueeze118_out1: Tensor<4, Bool>,
) -> Tensor<3> {
let layernormalization33_out1 = {
let dtype = add57_out1.clone().dtype();
self.layernormalization33
.forward(add57_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let transpose71_out1 = layernormalization33_out1.permute([0, 2, 1]);
let conv1d19_out1 = self.conv1d19.forward(transpose71_out1);
let split_tensors = conv1d19_out1.split_with_sizes([1024, 1024].into(), 1);
let [split7_out1, split7_out2] = split_tensors.try_into().unwrap();
let sigmoid26_out1 = burn::tensor::activation::sigmoid(split7_out2);
let mul85_out1 = split7_out1.mul(sigmoid26_out1);
let constant1654_out1 = 0f32;
let where39_out1 = mul85_out1.mask_fill(unsqueeze120_out1, constant1654_out1);
let pad14_out1 = where39_out1
.pad(
[(0usize, 0usize), (0usize, 0usize), (4usize, 4usize)],
burn::tensor::ops::PadMode::Constant(0f32),
);
let conv1d20_out1 = self.conv1d20.forward(pad14_out1);
let sigmoid27_out1 = burn::tensor::activation::sigmoid(conv1d20_out1.clone());
let mul86_out1 = conv1d20_out1.mul(sigmoid27_out1);
let conv1d21_out1 = self.conv1d21.forward(mul86_out1);
let transpose73_out1 = conv1d21_out1.permute([0, 2, 1]);
let add58_out1 = add57_out1.add(transpose73_out1);
let layernormalization34_out1 = {
let dtype = add58_out1.clone().dtype();
self.layernormalization34
.forward(add58_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let linear63_out1 = self.linear63.forward(layernormalization34_out1);
let sigmoid28_out1 = burn::tensor::activation::sigmoid(linear63_out1.clone());
let mul87_out1 = linear63_out1.mul(sigmoid28_out1);
let linear64_out1 = self.linear64.forward(mul87_out1);
let constant1663_out1 = self.constant1663.val();
let mul88_out1 = linear64_out1
.mul((constant1663_out1).unsqueeze_dims(&[0isize, 1isize]));
let add59_out1 = add58_out1.add(mul88_out1);
let layernormalization35_out1 = {
let dtype = add59_out1.dtype();
self.layernormalization35
.forward(add59_out1.cast(burn::tensor::DType::F32))
.cast(dtype)
};
let layernormalization36_out1 = {
let dtype = layernormalization35_out1.clone().dtype();
self.layernormalization36
.forward(
layernormalization35_out1.clone().cast(burn::tensor::DType::F32),
)
.cast(dtype)
};
let linear65_out1 = self.linear65.forward(layernormalization36_out1);
let sigmoid29_out1 = burn::tensor::activation::sigmoid(linear65_out1.clone());
let mul89_out1 = linear65_out1.mul(sigmoid29_out1);
let linear66_out1 = self.linear66.forward(mul89_out1);
let constant1664_out1 = self.constant1664.val();
let mul90_out1 = linear66_out1
.mul((constant1664_out1).unsqueeze_dims(&[0isize, 1isize]));
let add60_out1 = layernormalization35_out1.add(mul90_out1);
let layernormalization37_out1 = {
let dtype = add60_out1.clone().dtype();
self.layernormalization37
.forward(add60_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let shape111_out1: [i64; 3] = {
let axes = &layernormalization37_out1.clone().dims()[0..3];
let mut output = [0i64; 3];
for i in 0..3 {
output[i] = axes[i] as i64;
}
output
};
let gather85_out1 = shape111_out1[0] as i64;
let linear67_out1 = self.linear67.forward(layernormalization37_out1.clone());
let unsqueeze193_out1 = [gather85_out1 as i64];
let constant1669_out1: [i64; 1] = [128i64];
let constant1667_out1: [i64; 1] = [-1i64];
let constant1668_out1: [i64; 1] = [8i64];
let concat103_out1: [i64; 4usize] = [
&unsqueeze193_out1[..],
&constant1667_out1[..],
&constant1668_out1[..],
&constant1669_out1[..],
]
.concat()
.try_into()
.unwrap();
let unsqueeze194_out1 = [gather85_out1 as i64];
let constant1673_out1: [i64; 1] = [128i64];
let constant1671_out1: [i64; 1] = [-1i64];
let constant1672_out1: [i64; 1] = [8i64];
let concat104_out1: [i64; 4usize] = [
&unsqueeze194_out1[..],
&constant1671_out1[..],
&constant1672_out1[..],
&constant1673_out1[..],
]
.concat()
.try_into()
.unwrap();
let unsqueeze195_out1 = [gather85_out1 as i64];
let constant1677_out1: [i64; 1] = [128i64];
let constant1675_out1: [i64; 1] = [-1i64];
let constant1676_out1: [i64; 1] = [8i64];
let concat105_out1: [i64; 4usize] = [
&unsqueeze195_out1[..],
&constant1675_out1[..],
&constant1676_out1[..],
&constant1677_out1[..],
]
.concat()
.try_into()
.unwrap();
let reshape97_out1 = linear67_out1.reshape(concat103_out1);
let linear68_out1 = self.linear68.forward(layernormalization37_out1.clone());
let reshape98_out1 = linear68_out1.reshape(concat104_out1);
let linear69_out1 = self.linear69.forward(layernormalization37_out1);
let reshape99_out1 = linear69_out1.reshape(concat105_out1);
let transpose74_out1 = reshape99_out1.permute([0, 2, 1, 3]);
let linear70_out1 = self.linear70.forward(slice1_out1);
let reshape100_out1 = linear70_out1.reshape([1, -1, 8, 128]);
let constant192_out1 = self.constant192.val();
let add61_out1 = reshape97_out1
.clone()
.add((constant192_out1).unsqueeze_dims(&[0isize, 1isize]));
let transpose75_out1 = add61_out1.permute([0, 2, 1, 3]);
let constant193_out1 = self.constant193.val();
let add62_out1 = reshape97_out1
.add((constant193_out1).unsqueeze_dims(&[0isize, 1isize]));
let transpose76_out1 = add62_out1.permute([0, 2, 1, 3]);
let transpose77_out1 = reshape100_out1.permute([0, 2, 3, 1]);
let matmul92_out1 = transpose76_out1.matmul(transpose77_out1);
let shape112_out1: [i64; 4] = {
let axes = &matmul92_out1.clone().dims()[0..4];
let mut output = [0i64; 4];
for i in 0..4 {
output[i] = axes[i] as i64;
}
output
};
let gather86_out1 = shape112_out1[0] as i64;
let gather87_out1 = shape112_out1[1] as i64;
let gather88_out1 = shape112_out1[2] as i64;
let gather89_out1 = shape112_out1[3] as i64;
let pad15_out1 = matmul92_out1
.pad(
[(0usize, 0usize), (0usize, 0usize), (0usize, 0usize), (1usize, 0usize)],
burn::tensor::ops::PadMode::Constant(0f32),
);
let unsqueeze196_out1 = [gather86_out1 as i64];
let unsqueeze197_out1 = [gather87_out1 as i64];
let unsqueeze198_out1 = [gather88_out1 as i64];
let constant1692_out1: [i64; 1] = [-1i64];
let concat107_out1: [i64; 4usize] = [
&unsqueeze196_out1[..],
&unsqueeze197_out1[..],
&constant1692_out1[..],
&unsqueeze198_out1[..],
]
.concat()
.try_into()
.unwrap();
let reshape103_out1 = pad15_out1.reshape(concat107_out1);
let slice40_out1 = reshape103_out1.slice(s![.., .., 1.., ..]);
let unsqueeze199_out1 = [gather86_out1 as i64];
let unsqueeze200_out1 = [gather87_out1 as i64];
let unsqueeze201_out1 = [gather88_out1 as i64];
let unsqueeze202_out1 = [gather89_out1 as i64];
let concat108_out1: [i64; 4usize] = [
&unsqueeze199_out1[..],
&unsqueeze200_out1[..],
&unsqueeze201_out1[..],
&unsqueeze202_out1[..],
]
.concat()
.try_into()
.unwrap();
let reshape104_out1 = slice40_out1.reshape(concat108_out1);
let transpose79_out1 = reshape98_out1.permute([0, 2, 3, 1]);
let matmul93_out1 = transpose75_out1.matmul(transpose79_out1);
let shape116_out1: [i64; 4] = {
let axes = &matmul93_out1.clone().dims()[0..4];
let mut output = [0i64; 4];
for i in 0..4 {
output[i] = axes[i] as i64;
}
output
};
let slice41_out1: [i64; 1] = shape116_out1[3..4].try_into().unwrap();
let squeeze8_out1 = slice41_out1[0] as i64;
let unsqueeze203_out1 = [squeeze8_out1 as i64];
let slice42_out1 = reshape104_out1
.slice(s![.., .., .., 0..unsqueeze203_out1[0]]);
let add63_out1 = matmul93_out1.add(slice42_out1);
let constant1710_out1 = self.constant1710.val();
let div11_out1 = add63_out1
.div((constant1710_out1).unsqueeze_dims(&[0isize, 1isize, 2isize]));
let shape117_out1: [i64; 4] = {
let axes = &transpose74_out1.clone().dims()[0..4];
let mut output = [0i64; 4];
for i in 0..4 {
output[i] = axes[i] as i64;
}
output
};
let gather90_out1 = shape117_out1[0] as i64;
let constant1712_out1 = -10000f32;
let where40_out1 = div11_out1
.mask_fill(unsqueeze118_out1.clone(), constant1712_out1);
let softmax8_out1 = burn::tensor::activation::softmax(where40_out1, 3);
let constant1713_out1 = 0f32;
let where41_out1 = softmax8_out1.mask_fill(unsqueeze118_out1, constant1713_out1);
let matmul94_out1 = where41_out1.matmul(transpose74_out1);
let transpose80_out1 = matmul94_out1.permute([0, 2, 1, 3]);
let unsqueeze204_out1 = [gather90_out1 as i64];
let constant1716_out1: [i64; 1] = [1024i64];
let constant1715_out1: [i64; 1] = [-1i64];
let concat109_out1: [i64; 3usize] = [
&unsqueeze204_out1[..],
&constant1715_out1[..],
&constant1716_out1[..],
]
.concat()
.try_into()
.unwrap();
let reshape105_out1 = transpose80_out1.reshape(concat109_out1);
let linear71_out1 = self.linear71.forward(reshape105_out1);
let add64_out1 = add60_out1.add(linear71_out1);
add64_out1
}
}
#[derive(Module, Debug)]
pub struct Submodule13 {
layernormalization38: LayerNorm,
conv1d22: Conv1d,
conv1d23: Conv1d,
conv1d24: Conv1d,
layernormalization39: LayerNorm,
linear72: Linear,
linear73: Linear,
constant1726: burn::module::Param<Tensor<1>>,
layernormalization40: LayerNorm,
layernormalization41: LayerNorm,
linear74: Linear,
linear75: Linear,
constant1727: burn::module::Param<Tensor<1>>,
layernormalization42: LayerNorm,
linear76: Linear,
linear77: Linear,
linear78: Linear,
linear79: Linear,
constant216: burn::module::Param<Tensor<2>>,
constant217: burn::module::Param<Tensor<2>>,
constant1773: burn::module::Param<Tensor<1>>,
linear80: Linear,
#[module(skip)]
device: Device,
}
impl Submodule13 {
#[allow(unused_variables)]
pub fn new(device: &Device) -> Self {
let layernormalization38 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let conv1d22 = Conv1dConfig::new(1024, 2048, 1)
.with_stride(1)
.with_padding(PaddingConfig1d::Valid)
.with_dilation(1)
.with_groups(1)
.with_bias(true)
.init(device);
let conv1d23 = Conv1dConfig::new(1024, 1024, 9)
.with_stride(1)
.with_padding(PaddingConfig1d::Valid)
.with_dilation(1)
.with_groups(1024)
.with_bias(true)
.init(device);
let conv1d24 = Conv1dConfig::new(1024, 1024, 1)
.with_stride(1)
.with_padding(PaddingConfig1d::Valid)
.with_dilation(1)
.with_groups(1)
.with_bias(true)
.init(device);
let layernormalization39 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let linear72 = LinearConfig::new(1024, 4096).with_bias(true).init(device);
let linear73 = LinearConfig::new(4096, 1024).with_bias(true).init(device);
let constant1726: 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([0.5f64]),
(device, burn::tensor::DType::F32),
),
device.clone(),
false,
[1].into(),
);
let layernormalization40 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let layernormalization41 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let linear74 = LinearConfig::new(1024, 4096).with_bias(true).init(device);
let linear75 = LinearConfig::new(4096, 1024).with_bias(true).init(device);
let constant1727: 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([0.5f64]),
(device, burn::tensor::DType::F32),
),
device.clone(),
false,
[1].into(),
);
let layernormalization42 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let linear76 = LinearConfig::new(1024, 1024).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 linear79 = LinearConfig::new(1024, 1024).with_bias(false).init(device);
let constant216: burn::module::Param<Tensor<2>> = burn::module::Param::uninitialized(
burn::module::ParamId::new(),
move |device, _require_grad| Tensor::<
2,
>::zeros([8, 128], (device, burn::tensor::DType::F32)),
device.clone(),
false,
[8, 128].into(),
);
let constant217: burn::module::Param<Tensor<2>> = burn::module::Param::uninitialized(
burn::module::ParamId::new(),
move |device, _require_grad| Tensor::<
2,
>::zeros([8, 128], (device, burn::tensor::DType::F32)),
device.clone(),
false,
[8, 128].into(),
);
let constant1773: 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([11.313708305358887f64]),
(device, burn::tensor::DType::F32),
),
device.clone(),
false,
[1].into(),
);
let linear80 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
Self {
layernormalization38,
conv1d22,
conv1d23,
conv1d24,
layernormalization39,
linear72,
linear73,
constant1726,
layernormalization40,
layernormalization41,
linear74,
linear75,
constant1727,
layernormalization42,
linear76,
linear77,
linear78,
linear79,
constant216,
constant217,
constant1773,
linear80,
device: device.clone(),
}
}
#[allow(clippy::let_and_return, clippy::approx_constant)]
pub fn forward(
&self,
add64_out1: Tensor<3>,
unsqueeze120_out1: Tensor<3, Bool>,
slice1_out1: Tensor<3>,
unsqueeze118_out1: Tensor<4, Bool>,
) -> Tensor<3> {
let layernormalization38_out1 = {
let dtype = add64_out1.clone().dtype();
self.layernormalization38
.forward(add64_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let transpose81_out1 = layernormalization38_out1.permute([0, 2, 1]);
let conv1d22_out1 = self.conv1d22.forward(transpose81_out1);
let split_tensors = conv1d22_out1.split_with_sizes([1024, 1024].into(), 1);
let [split8_out1, split8_out2] = split_tensors.try_into().unwrap();
let sigmoid30_out1 = burn::tensor::activation::sigmoid(split8_out2);
let mul91_out1 = split8_out1.mul(sigmoid30_out1);
let constant1717_out1 = 0f32;
let where42_out1 = mul91_out1.mask_fill(unsqueeze120_out1, constant1717_out1);
let pad16_out1 = where42_out1
.pad(
[(0usize, 0usize), (0usize, 0usize), (4usize, 4usize)],
burn::tensor::ops::PadMode::Constant(0f32),
);
let conv1d23_out1 = self.conv1d23.forward(pad16_out1);
let sigmoid31_out1 = burn::tensor::activation::sigmoid(conv1d23_out1.clone());
let mul92_out1 = conv1d23_out1.mul(sigmoid31_out1);
let conv1d24_out1 = self.conv1d24.forward(mul92_out1);
let transpose83_out1 = conv1d24_out1.permute([0, 2, 1]);
let add65_out1 = add64_out1.add(transpose83_out1);
let layernormalization39_out1 = {
let dtype = add65_out1.clone().dtype();
self.layernormalization39
.forward(add65_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let linear72_out1 = self.linear72.forward(layernormalization39_out1);
let sigmoid32_out1 = burn::tensor::activation::sigmoid(linear72_out1.clone());
let mul93_out1 = linear72_out1.mul(sigmoid32_out1);
let linear73_out1 = self.linear73.forward(mul93_out1);
let constant1726_out1 = self.constant1726.val();
let mul94_out1 = linear73_out1
.mul((constant1726_out1).unsqueeze_dims(&[0isize, 1isize]));
let add66_out1 = add65_out1.add(mul94_out1);
let layernormalization40_out1 = {
let dtype = add66_out1.dtype();
self.layernormalization40
.forward(add66_out1.cast(burn::tensor::DType::F32))
.cast(dtype)
};
let layernormalization41_out1 = {
let dtype = layernormalization40_out1.clone().dtype();
self.layernormalization41
.forward(
layernormalization40_out1.clone().cast(burn::tensor::DType::F32),
)
.cast(dtype)
};
let linear74_out1 = self.linear74.forward(layernormalization41_out1);
let sigmoid33_out1 = burn::tensor::activation::sigmoid(linear74_out1.clone());
let mul95_out1 = linear74_out1.mul(sigmoid33_out1);
let linear75_out1 = self.linear75.forward(mul95_out1);
let constant1727_out1 = self.constant1727.val();
let mul96_out1 = linear75_out1
.mul((constant1727_out1).unsqueeze_dims(&[0isize, 1isize]));
let add67_out1 = layernormalization40_out1.add(mul96_out1);
let layernormalization42_out1 = {
let dtype = add67_out1.clone().dtype();
self.layernormalization42
.forward(add67_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let shape118_out1: [i64; 3] = {
let axes = &layernormalization42_out1.clone().dims()[0..3];
let mut output = [0i64; 3];
for i in 0..3 {
output[i] = axes[i] as i64;
}
output
};
let gather91_out1 = shape118_out1[0] as i64;
let linear76_out1 = self.linear76.forward(layernormalization42_out1.clone());
let unsqueeze205_out1 = [gather91_out1 as i64];
let constant1732_out1: [i64; 1] = [128i64];
let constant1730_out1: [i64; 1] = [-1i64];
let constant1731_out1: [i64; 1] = [8i64];
let concat111_out1: [i64; 4usize] = [
&unsqueeze205_out1[..],
&constant1730_out1[..],
&constant1731_out1[..],
&constant1732_out1[..],
]
.concat()
.try_into()
.unwrap();
let unsqueeze206_out1 = [gather91_out1 as i64];
let constant1736_out1: [i64; 1] = [128i64];
let constant1734_out1: [i64; 1] = [-1i64];
let constant1735_out1: [i64; 1] = [8i64];
let concat112_out1: [i64; 4usize] = [
&unsqueeze206_out1[..],
&constant1734_out1[..],
&constant1735_out1[..],
&constant1736_out1[..],
]
.concat()
.try_into()
.unwrap();
let unsqueeze207_out1 = [gather91_out1 as i64];
let constant1740_out1: [i64; 1] = [128i64];
let constant1738_out1: [i64; 1] = [-1i64];
let constant1739_out1: [i64; 1] = [8i64];
let concat113_out1: [i64; 4usize] = [
&unsqueeze207_out1[..],
&constant1738_out1[..],
&constant1739_out1[..],
&constant1740_out1[..],
]
.concat()
.try_into()
.unwrap();
let reshape108_out1 = linear76_out1.reshape(concat111_out1);
let linear77_out1 = self.linear77.forward(layernormalization42_out1.clone());
let reshape109_out1 = linear77_out1.reshape(concat112_out1);
let linear78_out1 = self.linear78.forward(layernormalization42_out1);
let reshape110_out1 = linear78_out1.reshape(concat113_out1);
let transpose84_out1 = reshape110_out1.permute([0, 2, 1, 3]);
let linear79_out1 = self.linear79.forward(slice1_out1);
let reshape111_out1 = linear79_out1.reshape([1, -1, 8, 128]);
let constant216_out1 = self.constant216.val();
let add68_out1 = reshape108_out1
.clone()
.add((constant216_out1).unsqueeze_dims(&[0isize, 1isize]));
let transpose85_out1 = add68_out1.permute([0, 2, 1, 3]);
let constant217_out1 = self.constant217.val();
let add69_out1 = reshape108_out1
.add((constant217_out1).unsqueeze_dims(&[0isize, 1isize]));
let transpose86_out1 = add69_out1.permute([0, 2, 1, 3]);
let transpose87_out1 = reshape111_out1.permute([0, 2, 3, 1]);
let matmul104_out1 = transpose86_out1.matmul(transpose87_out1);
let shape119_out1: [i64; 4] = {
let axes = &matmul104_out1.clone().dims()[0..4];
let mut output = [0i64; 4];
for i in 0..4 {
output[i] = axes[i] as i64;
}
output
};
let gather92_out1 = shape119_out1[0] as i64;
let gather93_out1 = shape119_out1[1] as i64;
let gather94_out1 = shape119_out1[2] as i64;
let gather95_out1 = shape119_out1[3] as i64;
let pad17_out1 = matmul104_out1
.pad(
[(0usize, 0usize), (0usize, 0usize), (0usize, 0usize), (1usize, 0usize)],
burn::tensor::ops::PadMode::Constant(0f32),
);
let unsqueeze208_out1 = [gather92_out1 as i64];
let unsqueeze209_out1 = [gather93_out1 as i64];
let unsqueeze210_out1 = [gather94_out1 as i64];
let constant1755_out1: [i64; 1] = [-1i64];
let concat115_out1: [i64; 4usize] = [
&unsqueeze208_out1[..],
&unsqueeze209_out1[..],
&constant1755_out1[..],
&unsqueeze210_out1[..],
]
.concat()
.try_into()
.unwrap();
let reshape114_out1 = pad17_out1.reshape(concat115_out1);
let slice45_out1 = reshape114_out1.slice(s![.., .., 1.., ..]);
let unsqueeze211_out1 = [gather92_out1 as i64];
let unsqueeze212_out1 = [gather93_out1 as i64];
let unsqueeze213_out1 = [gather94_out1 as i64];
let unsqueeze214_out1 = [gather95_out1 as i64];
let concat116_out1: [i64; 4usize] = [
&unsqueeze211_out1[..],
&unsqueeze212_out1[..],
&unsqueeze213_out1[..],
&unsqueeze214_out1[..],
]
.concat()
.try_into()
.unwrap();
let reshape115_out1 = slice45_out1.reshape(concat116_out1);
let transpose89_out1 = reshape109_out1.permute([0, 2, 3, 1]);
let matmul105_out1 = transpose85_out1.matmul(transpose89_out1);
let shape123_out1: [i64; 4] = {
let axes = &matmul105_out1.clone().dims()[0..4];
let mut output = [0i64; 4];
for i in 0..4 {
output[i] = axes[i] as i64;
}
output
};
let slice46_out1: [i64; 1] = shape123_out1[3..4].try_into().unwrap();
let squeeze9_out1 = slice46_out1[0] as i64;
let unsqueeze215_out1 = [squeeze9_out1 as i64];
let slice47_out1 = reshape115_out1
.slice(s![.., .., .., 0..unsqueeze215_out1[0]]);
let add70_out1 = matmul105_out1.add(slice47_out1);
let constant1773_out1 = self.constant1773.val();
let div12_out1 = add70_out1
.div((constant1773_out1).unsqueeze_dims(&[0isize, 1isize, 2isize]));
let shape124_out1: [i64; 4] = {
let axes = &transpose84_out1.clone().dims()[0..4];
let mut output = [0i64; 4];
for i in 0..4 {
output[i] = axes[i] as i64;
}
output
};
let gather96_out1 = shape124_out1[0] as i64;
let constant1775_out1 = -10000f32;
let where43_out1 = div12_out1
.mask_fill(unsqueeze118_out1.clone(), constant1775_out1);
let softmax9_out1 = burn::tensor::activation::softmax(where43_out1, 3);
let constant1776_out1 = 0f32;
let where44_out1 = softmax9_out1.mask_fill(unsqueeze118_out1, constant1776_out1);
let matmul106_out1 = where44_out1.matmul(transpose84_out1);
let transpose90_out1 = matmul106_out1.permute([0, 2, 1, 3]);
let unsqueeze216_out1 = [gather96_out1 as i64];
let constant1779_out1: [i64; 1] = [1024i64];
let constant1778_out1: [i64; 1] = [-1i64];
let concat117_out1: [i64; 3usize] = [
&unsqueeze216_out1[..],
&constant1778_out1[..],
&constant1779_out1[..],
]
.concat()
.try_into()
.unwrap();
let reshape116_out1 = transpose90_out1.reshape(concat117_out1);
let linear80_out1 = self.linear80.forward(reshape116_out1);
let add71_out1 = add67_out1.add(linear80_out1);
add71_out1
}
}
#[derive(Module, Debug)]
pub struct Submodule14 {
layernormalization43: LayerNorm,
conv1d25: Conv1d,
conv1d26: Conv1d,
conv1d27: Conv1d,
layernormalization44: LayerNorm,
linear81: Linear,
linear82: Linear,
constant1789: burn::module::Param<Tensor<1>>,
layernormalization45: LayerNorm,
layernormalization46: LayerNorm,
linear83: Linear,
linear84: Linear,
constant1790: burn::module::Param<Tensor<1>>,
layernormalization47: LayerNorm,
linear85: Linear,
linear86: Linear,
linear87: Linear,
linear88: Linear,
constant240: burn::module::Param<Tensor<2>>,
constant241: burn::module::Param<Tensor<2>>,
constant1836: burn::module::Param<Tensor<1>>,
linear89: Linear,
#[module(skip)]
device: Device,
}
impl Submodule14 {
#[allow(unused_variables)]
pub fn new(device: &Device) -> Self {
let layernormalization43 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let conv1d25 = Conv1dConfig::new(1024, 2048, 1)
.with_stride(1)
.with_padding(PaddingConfig1d::Valid)
.with_dilation(1)
.with_groups(1)
.with_bias(true)
.init(device);
let conv1d26 = Conv1dConfig::new(1024, 1024, 9)
.with_stride(1)
.with_padding(PaddingConfig1d::Valid)
.with_dilation(1)
.with_groups(1024)
.with_bias(true)
.init(device);
let conv1d27 = Conv1dConfig::new(1024, 1024, 1)
.with_stride(1)
.with_padding(PaddingConfig1d::Valid)
.with_dilation(1)
.with_groups(1)
.with_bias(true)
.init(device);
let layernormalization44 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let linear81 = LinearConfig::new(1024, 4096).with_bias(true).init(device);
let linear82 = LinearConfig::new(4096, 1024).with_bias(true).init(device);
let constant1789: 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([0.5f64]),
(device, burn::tensor::DType::F32),
),
device.clone(),
false,
[1].into(),
);
let layernormalization45 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let layernormalization46 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let linear83 = LinearConfig::new(1024, 4096).with_bias(true).init(device);
let linear84 = LinearConfig::new(4096, 1024).with_bias(true).init(device);
let constant1790: 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([0.5f64]),
(device, burn::tensor::DType::F32),
),
device.clone(),
false,
[1].into(),
);
let layernormalization47 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let linear85 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
let linear86 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
let linear87 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
let linear88 = LinearConfig::new(1024, 1024).with_bias(false).init(device);
let constant240: burn::module::Param<Tensor<2>> = burn::module::Param::uninitialized(
burn::module::ParamId::new(),
move |device, _require_grad| Tensor::<
2,
>::zeros([8, 128], (device, burn::tensor::DType::F32)),
device.clone(),
false,
[8, 128].into(),
);
let constant241: burn::module::Param<Tensor<2>> = burn::module::Param::uninitialized(
burn::module::ParamId::new(),
move |device, _require_grad| Tensor::<
2,
>::zeros([8, 128], (device, burn::tensor::DType::F32)),
device.clone(),
false,
[8, 128].into(),
);
let constant1836: 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([11.313708305358887f64]),
(device, burn::tensor::DType::F32),
),
device.clone(),
false,
[1].into(),
);
let linear89 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
Self {
layernormalization43,
conv1d25,
conv1d26,
conv1d27,
layernormalization44,
linear81,
linear82,
constant1789,
layernormalization45,
layernormalization46,
linear83,
linear84,
constant1790,
layernormalization47,
linear85,
linear86,
linear87,
linear88,
constant240,
constant241,
constant1836,
linear89,
device: device.clone(),
}
}
#[allow(clippy::let_and_return, clippy::approx_constant)]
pub fn forward(
&self,
add71_out1: Tensor<3>,
unsqueeze120_out1: Tensor<3, Bool>,
slice1_out1: Tensor<3>,
unsqueeze118_out1: Tensor<4, Bool>,
) -> Tensor<3> {
let layernormalization43_out1 = {
let dtype = add71_out1.clone().dtype();
self.layernormalization43
.forward(add71_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let transpose91_out1 = layernormalization43_out1.permute([0, 2, 1]);
let conv1d25_out1 = self.conv1d25.forward(transpose91_out1);
let split_tensors = conv1d25_out1.split_with_sizes([1024, 1024].into(), 1);
let [split9_out1, split9_out2] = split_tensors.try_into().unwrap();
let sigmoid34_out1 = burn::tensor::activation::sigmoid(split9_out2);
let mul97_out1 = split9_out1.mul(sigmoid34_out1);
let constant1780_out1 = 0f32;
let where45_out1 = mul97_out1.mask_fill(unsqueeze120_out1, constant1780_out1);
let pad18_out1 = where45_out1
.pad(
[(0usize, 0usize), (0usize, 0usize), (4usize, 4usize)],
burn::tensor::ops::PadMode::Constant(0f32),
);
let conv1d26_out1 = self.conv1d26.forward(pad18_out1);
let sigmoid35_out1 = burn::tensor::activation::sigmoid(conv1d26_out1.clone());
let mul98_out1 = conv1d26_out1.mul(sigmoid35_out1);
let conv1d27_out1 = self.conv1d27.forward(mul98_out1);
let transpose93_out1 = conv1d27_out1.permute([0, 2, 1]);
let add72_out1 = add71_out1.add(transpose93_out1);
let layernormalization44_out1 = {
let dtype = add72_out1.clone().dtype();
self.layernormalization44
.forward(add72_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let linear81_out1 = self.linear81.forward(layernormalization44_out1);
let sigmoid36_out1 = burn::tensor::activation::sigmoid(linear81_out1.clone());
let mul99_out1 = linear81_out1.mul(sigmoid36_out1);
let linear82_out1 = self.linear82.forward(mul99_out1);
let constant1789_out1 = self.constant1789.val();
let mul100_out1 = linear82_out1
.mul((constant1789_out1).unsqueeze_dims(&[0isize, 1isize]));
let add73_out1 = add72_out1.add(mul100_out1);
let layernormalization45_out1 = {
let dtype = add73_out1.dtype();
self.layernormalization45
.forward(add73_out1.cast(burn::tensor::DType::F32))
.cast(dtype)
};
let layernormalization46_out1 = {
let dtype = layernormalization45_out1.clone().dtype();
self.layernormalization46
.forward(
layernormalization45_out1.clone().cast(burn::tensor::DType::F32),
)
.cast(dtype)
};
let linear83_out1 = self.linear83.forward(layernormalization46_out1);
let sigmoid37_out1 = burn::tensor::activation::sigmoid(linear83_out1.clone());
let mul101_out1 = linear83_out1.mul(sigmoid37_out1);
let linear84_out1 = self.linear84.forward(mul101_out1);
let constant1790_out1 = self.constant1790.val();
let mul102_out1 = linear84_out1
.mul((constant1790_out1).unsqueeze_dims(&[0isize, 1isize]));
let add74_out1 = layernormalization45_out1.add(mul102_out1);
let layernormalization47_out1 = {
let dtype = add74_out1.clone().dtype();
self.layernormalization47
.forward(add74_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let shape125_out1: [i64; 3] = {
let axes = &layernormalization47_out1.clone().dims()[0..3];
let mut output = [0i64; 3];
for i in 0..3 {
output[i] = axes[i] as i64;
}
output
};
let gather97_out1 = shape125_out1[0] as i64;
let linear85_out1 = self.linear85.forward(layernormalization47_out1.clone());
let unsqueeze217_out1 = [gather97_out1 as i64];
let constant1795_out1: [i64; 1] = [128i64];
let constant1793_out1: [i64; 1] = [-1i64];
let constant1794_out1: [i64; 1] = [8i64];
let concat119_out1: [i64; 4usize] = [
&unsqueeze217_out1[..],
&constant1793_out1[..],
&constant1794_out1[..],
&constant1795_out1[..],
]
.concat()
.try_into()
.unwrap();
let unsqueeze218_out1 = [gather97_out1 as i64];
let constant1799_out1: [i64; 1] = [128i64];
let constant1797_out1: [i64; 1] = [-1i64];
let constant1798_out1: [i64; 1] = [8i64];
let concat120_out1: [i64; 4usize] = [
&unsqueeze218_out1[..],
&constant1797_out1[..],
&constant1798_out1[..],
&constant1799_out1[..],
]
.concat()
.try_into()
.unwrap();
let unsqueeze219_out1 = [gather97_out1 as i64];
let constant1803_out1: [i64; 1] = [128i64];
let constant1801_out1: [i64; 1] = [-1i64];
let constant1802_out1: [i64; 1] = [8i64];
let concat121_out1: [i64; 4usize] = [
&unsqueeze219_out1[..],
&constant1801_out1[..],
&constant1802_out1[..],
&constant1803_out1[..],
]
.concat()
.try_into()
.unwrap();
let reshape119_out1 = linear85_out1.reshape(concat119_out1);
let linear86_out1 = self.linear86.forward(layernormalization47_out1.clone());
let reshape120_out1 = linear86_out1.reshape(concat120_out1);
let linear87_out1 = self.linear87.forward(layernormalization47_out1);
let reshape121_out1 = linear87_out1.reshape(concat121_out1);
let transpose94_out1 = reshape121_out1.permute([0, 2, 1, 3]);
let linear88_out1 = self.linear88.forward(slice1_out1);
let reshape122_out1 = linear88_out1.reshape([1, -1, 8, 128]);
let constant240_out1 = self.constant240.val();
let add75_out1 = reshape119_out1
.clone()
.add((constant240_out1).unsqueeze_dims(&[0isize, 1isize]));
let transpose95_out1 = add75_out1.permute([0, 2, 1, 3]);
let constant241_out1 = self.constant241.val();
let add76_out1 = reshape119_out1
.add((constant241_out1).unsqueeze_dims(&[0isize, 1isize]));
let transpose96_out1 = add76_out1.permute([0, 2, 1, 3]);
let transpose97_out1 = reshape122_out1.permute([0, 2, 3, 1]);
let matmul116_out1 = transpose96_out1.matmul(transpose97_out1);
let shape126_out1: [i64; 4] = {
let axes = &matmul116_out1.clone().dims()[0..4];
let mut output = [0i64; 4];
for i in 0..4 {
output[i] = axes[i] as i64;
}
output
};
let gather98_out1 = shape126_out1[0] as i64;
let gather99_out1 = shape126_out1[1] as i64;
let gather100_out1 = shape126_out1[2] as i64;
let gather101_out1 = shape126_out1[3] as i64;
let pad19_out1 = matmul116_out1
.pad(
[(0usize, 0usize), (0usize, 0usize), (0usize, 0usize), (1usize, 0usize)],
burn::tensor::ops::PadMode::Constant(0f32),
);
let unsqueeze220_out1 = [gather98_out1 as i64];
let unsqueeze221_out1 = [gather99_out1 as i64];
let unsqueeze222_out1 = [gather100_out1 as i64];
let constant1818_out1: [i64; 1] = [-1i64];
let concat123_out1: [i64; 4usize] = [
&unsqueeze220_out1[..],
&unsqueeze221_out1[..],
&constant1818_out1[..],
&unsqueeze222_out1[..],
]
.concat()
.try_into()
.unwrap();
let reshape125_out1 = pad19_out1.reshape(concat123_out1);
let slice50_out1 = reshape125_out1.slice(s![.., .., 1.., ..]);
let unsqueeze223_out1 = [gather98_out1 as i64];
let unsqueeze224_out1 = [gather99_out1 as i64];
let unsqueeze225_out1 = [gather100_out1 as i64];
let unsqueeze226_out1 = [gather101_out1 as i64];
let concat124_out1: [i64; 4usize] = [
&unsqueeze223_out1[..],
&unsqueeze224_out1[..],
&unsqueeze225_out1[..],
&unsqueeze226_out1[..],
]
.concat()
.try_into()
.unwrap();
let reshape126_out1 = slice50_out1.reshape(concat124_out1);
let transpose99_out1 = reshape120_out1.permute([0, 2, 3, 1]);
let matmul117_out1 = transpose95_out1.matmul(transpose99_out1);
let shape130_out1: [i64; 4] = {
let axes = &matmul117_out1.clone().dims()[0..4];
let mut output = [0i64; 4];
for i in 0..4 {
output[i] = axes[i] as i64;
}
output
};
let slice51_out1: [i64; 1] = shape130_out1[3..4].try_into().unwrap();
let squeeze10_out1 = slice51_out1[0] as i64;
let unsqueeze227_out1 = [squeeze10_out1 as i64];
let slice52_out1 = reshape126_out1
.slice(s![.., .., .., 0..unsqueeze227_out1[0]]);
let add77_out1 = matmul117_out1.add(slice52_out1);
let constant1836_out1 = self.constant1836.val();
let div13_out1 = add77_out1
.div((constant1836_out1).unsqueeze_dims(&[0isize, 1isize, 2isize]));
let shape131_out1: [i64; 4] = {
let axes = &transpose94_out1.clone().dims()[0..4];
let mut output = [0i64; 4];
for i in 0..4 {
output[i] = axes[i] as i64;
}
output
};
let gather102_out1 = shape131_out1[0] as i64;
let constant1838_out1 = -10000f32;
let where46_out1 = div13_out1
.mask_fill(unsqueeze118_out1.clone(), constant1838_out1);
let softmax10_out1 = burn::tensor::activation::softmax(where46_out1, 3);
let constant1839_out1 = 0f32;
let where47_out1 = softmax10_out1
.mask_fill(unsqueeze118_out1, constant1839_out1);
let matmul118_out1 = where47_out1.matmul(transpose94_out1);
let transpose100_out1 = matmul118_out1.permute([0, 2, 1, 3]);
let unsqueeze228_out1 = [gather102_out1 as i64];
let constant1842_out1: [i64; 1] = [1024i64];
let constant1841_out1: [i64; 1] = [-1i64];
let concat125_out1: [i64; 3usize] = [
&unsqueeze228_out1[..],
&constant1841_out1[..],
&constant1842_out1[..],
]
.concat()
.try_into()
.unwrap();
let reshape127_out1 = transpose100_out1.reshape(concat125_out1);
let linear89_out1 = self.linear89.forward(reshape127_out1);
let add78_out1 = add74_out1.add(linear89_out1);
add78_out1
}
}
#[derive(Module, Debug)]
pub struct Submodule15 {
layernormalization48: LayerNorm,
conv1d28: Conv1d,
conv1d29: Conv1d,
conv1d30: Conv1d,
layernormalization49: LayerNorm,
linear90: Linear,
linear91: Linear,
constant1852: burn::module::Param<Tensor<1>>,
layernormalization50: LayerNorm,
layernormalization51: LayerNorm,
linear92: Linear,
linear93: Linear,
constant1853: burn::module::Param<Tensor<1>>,
layernormalization52: LayerNorm,
linear94: Linear,
linear95: Linear,
linear96: Linear,
linear97: Linear,
constant264: burn::module::Param<Tensor<2>>,
constant265: burn::module::Param<Tensor<2>>,
constant1899: burn::module::Param<Tensor<1>>,
linear98: Linear,
#[module(skip)]
device: Device,
}
impl Submodule15 {
#[allow(unused_variables)]
pub fn new(device: &Device) -> Self {
let layernormalization48 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let conv1d28 = Conv1dConfig::new(1024, 2048, 1)
.with_stride(1)
.with_padding(PaddingConfig1d::Valid)
.with_dilation(1)
.with_groups(1)
.with_bias(true)
.init(device);
let conv1d29 = Conv1dConfig::new(1024, 1024, 9)
.with_stride(1)
.with_padding(PaddingConfig1d::Valid)
.with_dilation(1)
.with_groups(1024)
.with_bias(true)
.init(device);
let conv1d30 = Conv1dConfig::new(1024, 1024, 1)
.with_stride(1)
.with_padding(PaddingConfig1d::Valid)
.with_dilation(1)
.with_groups(1)
.with_bias(true)
.init(device);
let layernormalization49 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let linear90 = LinearConfig::new(1024, 4096).with_bias(true).init(device);
let linear91 = LinearConfig::new(4096, 1024).with_bias(true).init(device);
let constant1852: 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([0.5f64]),
(device, burn::tensor::DType::F32),
),
device.clone(),
false,
[1].into(),
);
let layernormalization50 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let layernormalization51 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let linear92 = LinearConfig::new(1024, 4096).with_bias(true).init(device);
let linear93 = LinearConfig::new(4096, 1024).with_bias(true).init(device);
let constant1853: 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([0.5f64]),
(device, burn::tensor::DType::F32),
),
device.clone(),
false,
[1].into(),
);
let layernormalization52 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let linear94 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
let linear95 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
let linear96 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
let linear97 = LinearConfig::new(1024, 1024).with_bias(false).init(device);
let constant264: burn::module::Param<Tensor<2>> = burn::module::Param::uninitialized(
burn::module::ParamId::new(),
move |device, _require_grad| Tensor::<
2,
>::zeros([8, 128], (device, burn::tensor::DType::F32)),
device.clone(),
false,
[8, 128].into(),
);
let constant265: burn::module::Param<Tensor<2>> = burn::module::Param::uninitialized(
burn::module::ParamId::new(),
move |device, _require_grad| Tensor::<
2,
>::zeros([8, 128], (device, burn::tensor::DType::F32)),
device.clone(),
false,
[8, 128].into(),
);
let constant1899: 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([11.313708305358887f64]),
(device, burn::tensor::DType::F32),
),
device.clone(),
false,
[1].into(),
);
let linear98 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
Self {
layernormalization48,
conv1d28,
conv1d29,
conv1d30,
layernormalization49,
linear90,
linear91,
constant1852,
layernormalization50,
layernormalization51,
linear92,
linear93,
constant1853,
layernormalization52,
linear94,
linear95,
linear96,
linear97,
constant264,
constant265,
constant1899,
linear98,
device: device.clone(),
}
}
#[allow(clippy::let_and_return, clippy::approx_constant)]
pub fn forward(
&self,
add78_out1: Tensor<3>,
unsqueeze120_out1: Tensor<3, Bool>,
slice1_out1: Tensor<3>,
unsqueeze118_out1: Tensor<4, Bool>,
) -> Tensor<3> {
let layernormalization48_out1 = {
let dtype = add78_out1.clone().dtype();
self.layernormalization48
.forward(add78_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let transpose101_out1 = layernormalization48_out1.permute([0, 2, 1]);
let conv1d28_out1 = self.conv1d28.forward(transpose101_out1);
let split_tensors = conv1d28_out1.split_with_sizes([1024, 1024].into(), 1);
let [split10_out1, split10_out2] = split_tensors.try_into().unwrap();
let sigmoid38_out1 = burn::tensor::activation::sigmoid(split10_out2);
let mul103_out1 = split10_out1.mul(sigmoid38_out1);
let constant1843_out1 = 0f32;
let where48_out1 = mul103_out1.mask_fill(unsqueeze120_out1, constant1843_out1);
let pad20_out1 = where48_out1
.pad(
[(0usize, 0usize), (0usize, 0usize), (4usize, 4usize)],
burn::tensor::ops::PadMode::Constant(0f32),
);
let conv1d29_out1 = self.conv1d29.forward(pad20_out1);
let sigmoid39_out1 = burn::tensor::activation::sigmoid(conv1d29_out1.clone());
let mul104_out1 = conv1d29_out1.mul(sigmoid39_out1);
let conv1d30_out1 = self.conv1d30.forward(mul104_out1);
let transpose103_out1 = conv1d30_out1.permute([0, 2, 1]);
let add79_out1 = add78_out1.add(transpose103_out1);
let layernormalization49_out1 = {
let dtype = add79_out1.clone().dtype();
self.layernormalization49
.forward(add79_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let linear90_out1 = self.linear90.forward(layernormalization49_out1);
let sigmoid40_out1 = burn::tensor::activation::sigmoid(linear90_out1.clone());
let mul105_out1 = linear90_out1.mul(sigmoid40_out1);
let linear91_out1 = self.linear91.forward(mul105_out1);
let constant1852_out1 = self.constant1852.val();
let mul106_out1 = linear91_out1
.mul((constant1852_out1).unsqueeze_dims(&[0isize, 1isize]));
let add80_out1 = add79_out1.add(mul106_out1);
let layernormalization50_out1 = {
let dtype = add80_out1.dtype();
self.layernormalization50
.forward(add80_out1.cast(burn::tensor::DType::F32))
.cast(dtype)
};
let layernormalization51_out1 = {
let dtype = layernormalization50_out1.clone().dtype();
self.layernormalization51
.forward(
layernormalization50_out1.clone().cast(burn::tensor::DType::F32),
)
.cast(dtype)
};
let linear92_out1 = self.linear92.forward(layernormalization51_out1);
let sigmoid41_out1 = burn::tensor::activation::sigmoid(linear92_out1.clone());
let mul107_out1 = linear92_out1.mul(sigmoid41_out1);
let linear93_out1 = self.linear93.forward(mul107_out1);
let constant1853_out1 = self.constant1853.val();
let mul108_out1 = linear93_out1
.mul((constant1853_out1).unsqueeze_dims(&[0isize, 1isize]));
let add81_out1 = layernormalization50_out1.add(mul108_out1);
let layernormalization52_out1 = {
let dtype = add81_out1.clone().dtype();
self.layernormalization52
.forward(add81_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let shape132_out1: [i64; 3] = {
let axes = &layernormalization52_out1.clone().dims()[0..3];
let mut output = [0i64; 3];
for i in 0..3 {
output[i] = axes[i] as i64;
}
output
};
let gather103_out1 = shape132_out1[0] as i64;
let linear94_out1 = self.linear94.forward(layernormalization52_out1.clone());
let unsqueeze229_out1 = [gather103_out1 as i64];
let constant1858_out1: [i64; 1] = [128i64];
let constant1856_out1: [i64; 1] = [-1i64];
let constant1857_out1: [i64; 1] = [8i64];
let concat127_out1: [i64; 4usize] = [
&unsqueeze229_out1[..],
&constant1856_out1[..],
&constant1857_out1[..],
&constant1858_out1[..],
]
.concat()
.try_into()
.unwrap();
let unsqueeze230_out1 = [gather103_out1 as i64];
let constant1862_out1: [i64; 1] = [128i64];
let constant1860_out1: [i64; 1] = [-1i64];
let constant1861_out1: [i64; 1] = [8i64];
let concat128_out1: [i64; 4usize] = [
&unsqueeze230_out1[..],
&constant1860_out1[..],
&constant1861_out1[..],
&constant1862_out1[..],
]
.concat()
.try_into()
.unwrap();
let unsqueeze231_out1 = [gather103_out1 as i64];
let constant1866_out1: [i64; 1] = [128i64];
let constant1864_out1: [i64; 1] = [-1i64];
let constant1865_out1: [i64; 1] = [8i64];
let concat129_out1: [i64; 4usize] = [
&unsqueeze231_out1[..],
&constant1864_out1[..],
&constant1865_out1[..],
&constant1866_out1[..],
]
.concat()
.try_into()
.unwrap();
let reshape130_out1 = linear94_out1.reshape(concat127_out1);
let linear95_out1 = self.linear95.forward(layernormalization52_out1.clone());
let reshape131_out1 = linear95_out1.reshape(concat128_out1);
let linear96_out1 = self.linear96.forward(layernormalization52_out1);
let reshape132_out1 = linear96_out1.reshape(concat129_out1);
let transpose104_out1 = reshape132_out1.permute([0, 2, 1, 3]);
let linear97_out1 = self.linear97.forward(slice1_out1);
let reshape133_out1 = linear97_out1.reshape([1, -1, 8, 128]);
let constant264_out1 = self.constant264.val();
let add82_out1 = reshape130_out1
.clone()
.add((constant264_out1).unsqueeze_dims(&[0isize, 1isize]));
let transpose105_out1 = add82_out1.permute([0, 2, 1, 3]);
let constant265_out1 = self.constant265.val();
let add83_out1 = reshape130_out1
.add((constant265_out1).unsqueeze_dims(&[0isize, 1isize]));
let transpose106_out1 = add83_out1.permute([0, 2, 1, 3]);
let transpose107_out1 = reshape133_out1.permute([0, 2, 3, 1]);
let matmul128_out1 = transpose106_out1.matmul(transpose107_out1);
let shape133_out1: [i64; 4] = {
let axes = &matmul128_out1.clone().dims()[0..4];
let mut output = [0i64; 4];
for i in 0..4 {
output[i] = axes[i] as i64;
}
output
};
let gather104_out1 = shape133_out1[0] as i64;
let gather105_out1 = shape133_out1[1] as i64;
let gather106_out1 = shape133_out1[2] as i64;
let gather107_out1 = shape133_out1[3] as i64;
let pad21_out1 = matmul128_out1
.pad(
[(0usize, 0usize), (0usize, 0usize), (0usize, 0usize), (1usize, 0usize)],
burn::tensor::ops::PadMode::Constant(0f32),
);
let unsqueeze232_out1 = [gather104_out1 as i64];
let unsqueeze233_out1 = [gather105_out1 as i64];
let unsqueeze234_out1 = [gather106_out1 as i64];
let constant1881_out1: [i64; 1] = [-1i64];
let concat131_out1: [i64; 4usize] = [
&unsqueeze232_out1[..],
&unsqueeze233_out1[..],
&constant1881_out1[..],
&unsqueeze234_out1[..],
]
.concat()
.try_into()
.unwrap();
let reshape136_out1 = pad21_out1.reshape(concat131_out1);
let slice55_out1 = reshape136_out1.slice(s![.., .., 1.., ..]);
let unsqueeze235_out1 = [gather104_out1 as i64];
let unsqueeze236_out1 = [gather105_out1 as i64];
let unsqueeze237_out1 = [gather106_out1 as i64];
let unsqueeze238_out1 = [gather107_out1 as i64];
let concat132_out1: [i64; 4usize] = [
&unsqueeze235_out1[..],
&unsqueeze236_out1[..],
&unsqueeze237_out1[..],
&unsqueeze238_out1[..],
]
.concat()
.try_into()
.unwrap();
let reshape137_out1 = slice55_out1.reshape(concat132_out1);
let transpose109_out1 = reshape131_out1.permute([0, 2, 3, 1]);
let matmul129_out1 = transpose105_out1.matmul(transpose109_out1);
let shape137_out1: [i64; 4] = {
let axes = &matmul129_out1.clone().dims()[0..4];
let mut output = [0i64; 4];
for i in 0..4 {
output[i] = axes[i] as i64;
}
output
};
let slice56_out1: [i64; 1] = shape137_out1[3..4].try_into().unwrap();
let squeeze11_out1 = slice56_out1[0] as i64;
let unsqueeze239_out1 = [squeeze11_out1 as i64];
let slice57_out1 = reshape137_out1
.slice(s![.., .., .., 0..unsqueeze239_out1[0]]);
let add84_out1 = matmul129_out1.add(slice57_out1);
let constant1899_out1 = self.constant1899.val();
let div14_out1 = add84_out1
.div((constant1899_out1).unsqueeze_dims(&[0isize, 1isize, 2isize]));
let shape138_out1: [i64; 4] = {
let axes = &transpose104_out1.clone().dims()[0..4];
let mut output = [0i64; 4];
for i in 0..4 {
output[i] = axes[i] as i64;
}
output
};
let gather108_out1 = shape138_out1[0] as i64;
let constant1901_out1 = -10000f32;
let where49_out1 = div14_out1
.mask_fill(unsqueeze118_out1.clone(), constant1901_out1);
let softmax11_out1 = burn::tensor::activation::softmax(where49_out1, 3);
let constant1902_out1 = 0f32;
let where50_out1 = softmax11_out1
.mask_fill(unsqueeze118_out1, constant1902_out1);
let matmul130_out1 = where50_out1.matmul(transpose104_out1);
let transpose110_out1 = matmul130_out1.permute([0, 2, 1, 3]);
let unsqueeze240_out1 = [gather108_out1 as i64];
let constant1905_out1: [i64; 1] = [1024i64];
let constant1904_out1: [i64; 1] = [-1i64];
let concat133_out1: [i64; 3usize] = [
&unsqueeze240_out1[..],
&constant1904_out1[..],
&constant1905_out1[..],
]
.concat()
.try_into()
.unwrap();
let reshape138_out1 = transpose110_out1.reshape(concat133_out1);
let linear98_out1 = self.linear98.forward(reshape138_out1);
let add85_out1 = add81_out1.add(linear98_out1);
add85_out1
}
}
#[derive(Module, Debug)]
pub struct Submodule16 {
layernormalization53: LayerNorm,
conv1d31: Conv1d,
conv1d32: Conv1d,
conv1d33: Conv1d,
layernormalization54: LayerNorm,
linear99: Linear,
linear100: Linear,
constant1915: burn::module::Param<Tensor<1>>,
layernormalization55: LayerNorm,
layernormalization56: LayerNorm,
linear101: Linear,
linear102: Linear,
constant1916: burn::module::Param<Tensor<1>>,
layernormalization57: LayerNorm,
linear103: Linear,
linear104: Linear,
linear105: Linear,
linear106: Linear,
constant288: burn::module::Param<Tensor<2>>,
constant289: burn::module::Param<Tensor<2>>,
constant1962: burn::module::Param<Tensor<1>>,
linear107: Linear,
#[module(skip)]
device: Device,
}
impl Submodule16 {
#[allow(unused_variables)]
pub fn new(device: &Device) -> Self {
let layernormalization53 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let conv1d31 = Conv1dConfig::new(1024, 2048, 1)
.with_stride(1)
.with_padding(PaddingConfig1d::Valid)
.with_dilation(1)
.with_groups(1)
.with_bias(true)
.init(device);
let conv1d32 = Conv1dConfig::new(1024, 1024, 9)
.with_stride(1)
.with_padding(PaddingConfig1d::Valid)
.with_dilation(1)
.with_groups(1024)
.with_bias(true)
.init(device);
let conv1d33 = Conv1dConfig::new(1024, 1024, 1)
.with_stride(1)
.with_padding(PaddingConfig1d::Valid)
.with_dilation(1)
.with_groups(1)
.with_bias(true)
.init(device);
let layernormalization54 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let linear99 = LinearConfig::new(1024, 4096).with_bias(true).init(device);
let linear100 = LinearConfig::new(4096, 1024).with_bias(true).init(device);
let constant1915: 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([0.5f64]),
(device, burn::tensor::DType::F32),
),
device.clone(),
false,
[1].into(),
);
let layernormalization55 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let layernormalization56 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let linear101 = LinearConfig::new(1024, 4096).with_bias(true).init(device);
let linear102 = LinearConfig::new(4096, 1024).with_bias(true).init(device);
let constant1916: 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([0.5f64]),
(device, burn::tensor::DType::F32),
),
device.clone(),
false,
[1].into(),
);
let layernormalization57 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let linear103 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
let linear104 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
let linear105 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
let linear106 = LinearConfig::new(1024, 1024).with_bias(false).init(device);
let constant288: burn::module::Param<Tensor<2>> = burn::module::Param::uninitialized(
burn::module::ParamId::new(),
move |device, _require_grad| Tensor::<
2,
>::zeros([8, 128], (device, burn::tensor::DType::F32)),
device.clone(),
false,
[8, 128].into(),
);
let constant289: burn::module::Param<Tensor<2>> = burn::module::Param::uninitialized(
burn::module::ParamId::new(),
move |device, _require_grad| Tensor::<
2,
>::zeros([8, 128], (device, burn::tensor::DType::F32)),
device.clone(),
false,
[8, 128].into(),
);
let constant1962: 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([11.313708305358887f64]),
(device, burn::tensor::DType::F32),
),
device.clone(),
false,
[1].into(),
);
let linear107 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
Self {
layernormalization53,
conv1d31,
conv1d32,
conv1d33,
layernormalization54,
linear99,
linear100,
constant1915,
layernormalization55,
layernormalization56,
linear101,
linear102,
constant1916,
layernormalization57,
linear103,
linear104,
linear105,
linear106,
constant288,
constant289,
constant1962,
linear107,
device: device.clone(),
}
}
#[allow(clippy::let_and_return, clippy::approx_constant)]
pub fn forward(
&self,
add85_out1: Tensor<3>,
unsqueeze120_out1: Tensor<3, Bool>,
slice1_out1: Tensor<3>,
unsqueeze118_out1: Tensor<4, Bool>,
) -> Tensor<3> {
let layernormalization53_out1 = {
let dtype = add85_out1.clone().dtype();
self.layernormalization53
.forward(add85_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let transpose111_out1 = layernormalization53_out1.permute([0, 2, 1]);
let conv1d31_out1 = self.conv1d31.forward(transpose111_out1);
let split_tensors = conv1d31_out1.split_with_sizes([1024, 1024].into(), 1);
let [split11_out1, split11_out2] = split_tensors.try_into().unwrap();
let sigmoid42_out1 = burn::tensor::activation::sigmoid(split11_out2);
let mul109_out1 = split11_out1.mul(sigmoid42_out1);
let constant1906_out1 = 0f32;
let where51_out1 = mul109_out1.mask_fill(unsqueeze120_out1, constant1906_out1);
let pad22_out1 = where51_out1
.pad(
[(0usize, 0usize), (0usize, 0usize), (4usize, 4usize)],
burn::tensor::ops::PadMode::Constant(0f32),
);
let conv1d32_out1 = self.conv1d32.forward(pad22_out1);
let sigmoid43_out1 = burn::tensor::activation::sigmoid(conv1d32_out1.clone());
let mul110_out1 = conv1d32_out1.mul(sigmoid43_out1);
let conv1d33_out1 = self.conv1d33.forward(mul110_out1);
let transpose113_out1 = conv1d33_out1.permute([0, 2, 1]);
let add86_out1 = add85_out1.add(transpose113_out1);
let layernormalization54_out1 = {
let dtype = add86_out1.clone().dtype();
self.layernormalization54
.forward(add86_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let linear99_out1 = self.linear99.forward(layernormalization54_out1);
let sigmoid44_out1 = burn::tensor::activation::sigmoid(linear99_out1.clone());
let mul111_out1 = linear99_out1.mul(sigmoid44_out1);
let linear100_out1 = self.linear100.forward(mul111_out1);
let constant1915_out1 = self.constant1915.val();
let mul112_out1 = linear100_out1
.mul((constant1915_out1).unsqueeze_dims(&[0isize, 1isize]));
let add87_out1 = add86_out1.add(mul112_out1);
let layernormalization55_out1 = {
let dtype = add87_out1.dtype();
self.layernormalization55
.forward(add87_out1.cast(burn::tensor::DType::F32))
.cast(dtype)
};
let layernormalization56_out1 = {
let dtype = layernormalization55_out1.clone().dtype();
self.layernormalization56
.forward(
layernormalization55_out1.clone().cast(burn::tensor::DType::F32),
)
.cast(dtype)
};
let linear101_out1 = self.linear101.forward(layernormalization56_out1);
let sigmoid45_out1 = burn::tensor::activation::sigmoid(linear101_out1.clone());
let mul113_out1 = linear101_out1.mul(sigmoid45_out1);
let linear102_out1 = self.linear102.forward(mul113_out1);
let constant1916_out1 = self.constant1916.val();
let mul114_out1 = linear102_out1
.mul((constant1916_out1).unsqueeze_dims(&[0isize, 1isize]));
let add88_out1 = layernormalization55_out1.add(mul114_out1);
let layernormalization57_out1 = {
let dtype = add88_out1.clone().dtype();
self.layernormalization57
.forward(add88_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let shape139_out1: [i64; 3] = {
let axes = &layernormalization57_out1.clone().dims()[0..3];
let mut output = [0i64; 3];
for i in 0..3 {
output[i] = axes[i] as i64;
}
output
};
let gather109_out1 = shape139_out1[0] as i64;
let linear103_out1 = self.linear103.forward(layernormalization57_out1.clone());
let unsqueeze241_out1 = [gather109_out1 as i64];
let constant1921_out1: [i64; 1] = [128i64];
let constant1919_out1: [i64; 1] = [-1i64];
let constant1920_out1: [i64; 1] = [8i64];
let concat135_out1: [i64; 4usize] = [
&unsqueeze241_out1[..],
&constant1919_out1[..],
&constant1920_out1[..],
&constant1921_out1[..],
]
.concat()
.try_into()
.unwrap();
let unsqueeze242_out1 = [gather109_out1 as i64];
let constant1925_out1: [i64; 1] = [128i64];
let constant1923_out1: [i64; 1] = [-1i64];
let constant1924_out1: [i64; 1] = [8i64];
let concat136_out1: [i64; 4usize] = [
&unsqueeze242_out1[..],
&constant1923_out1[..],
&constant1924_out1[..],
&constant1925_out1[..],
]
.concat()
.try_into()
.unwrap();
let unsqueeze243_out1 = [gather109_out1 as i64];
let constant1929_out1: [i64; 1] = [128i64];
let constant1927_out1: [i64; 1] = [-1i64];
let constant1928_out1: [i64; 1] = [8i64];
let concat137_out1: [i64; 4usize] = [
&unsqueeze243_out1[..],
&constant1927_out1[..],
&constant1928_out1[..],
&constant1929_out1[..],
]
.concat()
.try_into()
.unwrap();
let reshape141_out1 = linear103_out1.reshape(concat135_out1);
let linear104_out1 = self.linear104.forward(layernormalization57_out1.clone());
let reshape142_out1 = linear104_out1.reshape(concat136_out1);
let linear105_out1 = self.linear105.forward(layernormalization57_out1);
let reshape143_out1 = linear105_out1.reshape(concat137_out1);
let transpose114_out1 = reshape143_out1.permute([0, 2, 1, 3]);
let linear106_out1 = self.linear106.forward(slice1_out1);
let reshape144_out1 = linear106_out1.reshape([1, -1, 8, 128]);
let constant288_out1 = self.constant288.val();
let add89_out1 = reshape141_out1
.clone()
.add((constant288_out1).unsqueeze_dims(&[0isize, 1isize]));
let transpose115_out1 = add89_out1.permute([0, 2, 1, 3]);
let constant289_out1 = self.constant289.val();
let add90_out1 = reshape141_out1
.add((constant289_out1).unsqueeze_dims(&[0isize, 1isize]));
let transpose116_out1 = add90_out1.permute([0, 2, 1, 3]);
let transpose117_out1 = reshape144_out1.permute([0, 2, 3, 1]);
let matmul140_out1 = transpose116_out1.matmul(transpose117_out1);
let shape140_out1: [i64; 4] = {
let axes = &matmul140_out1.clone().dims()[0..4];
let mut output = [0i64; 4];
for i in 0..4 {
output[i] = axes[i] as i64;
}
output
};
let gather110_out1 = shape140_out1[0] as i64;
let gather111_out1 = shape140_out1[1] as i64;
let gather112_out1 = shape140_out1[2] as i64;
let gather113_out1 = shape140_out1[3] as i64;
let pad23_out1 = matmul140_out1
.pad(
[(0usize, 0usize), (0usize, 0usize), (0usize, 0usize), (1usize, 0usize)],
burn::tensor::ops::PadMode::Constant(0f32),
);
let unsqueeze244_out1 = [gather110_out1 as i64];
let unsqueeze245_out1 = [gather111_out1 as i64];
let unsqueeze246_out1 = [gather112_out1 as i64];
let constant1944_out1: [i64; 1] = [-1i64];
let concat139_out1: [i64; 4usize] = [
&unsqueeze244_out1[..],
&unsqueeze245_out1[..],
&constant1944_out1[..],
&unsqueeze246_out1[..],
]
.concat()
.try_into()
.unwrap();
let reshape147_out1 = pad23_out1.reshape(concat139_out1);
let slice60_out1 = reshape147_out1.slice(s![.., .., 1.., ..]);
let unsqueeze247_out1 = [gather110_out1 as i64];
let unsqueeze248_out1 = [gather111_out1 as i64];
let unsqueeze249_out1 = [gather112_out1 as i64];
let unsqueeze250_out1 = [gather113_out1 as i64];
let concat140_out1: [i64; 4usize] = [
&unsqueeze247_out1[..],
&unsqueeze248_out1[..],
&unsqueeze249_out1[..],
&unsqueeze250_out1[..],
]
.concat()
.try_into()
.unwrap();
let reshape148_out1 = slice60_out1.reshape(concat140_out1);
let transpose119_out1 = reshape142_out1.permute([0, 2, 3, 1]);
let matmul141_out1 = transpose115_out1.matmul(transpose119_out1);
let shape144_out1: [i64; 4] = {
let axes = &matmul141_out1.clone().dims()[0..4];
let mut output = [0i64; 4];
for i in 0..4 {
output[i] = axes[i] as i64;
}
output
};
let slice61_out1: [i64; 1] = shape144_out1[3..4].try_into().unwrap();
let squeeze12_out1 = slice61_out1[0] as i64;
let unsqueeze251_out1 = [squeeze12_out1 as i64];
let slice62_out1 = reshape148_out1
.slice(s![.., .., .., 0..unsqueeze251_out1[0]]);
let add91_out1 = matmul141_out1.add(slice62_out1);
let constant1962_out1 = self.constant1962.val();
let div15_out1 = add91_out1
.div((constant1962_out1).unsqueeze_dims(&[0isize, 1isize, 2isize]));
let shape145_out1: [i64; 4] = {
let axes = &transpose114_out1.clone().dims()[0..4];
let mut output = [0i64; 4];
for i in 0..4 {
output[i] = axes[i] as i64;
}
output
};
let gather114_out1 = shape145_out1[0] as i64;
let constant1964_out1 = -10000f32;
let where52_out1 = div15_out1
.mask_fill(unsqueeze118_out1.clone(), constant1964_out1);
let softmax12_out1 = burn::tensor::activation::softmax(where52_out1, 3);
let constant1965_out1 = 0f32;
let where53_out1 = softmax12_out1
.mask_fill(unsqueeze118_out1, constant1965_out1);
let matmul142_out1 = where53_out1.matmul(transpose114_out1);
let transpose120_out1 = matmul142_out1.permute([0, 2, 1, 3]);
let unsqueeze252_out1 = [gather114_out1 as i64];
let constant1968_out1: [i64; 1] = [1024i64];
let constant1967_out1: [i64; 1] = [-1i64];
let concat141_out1: [i64; 3usize] = [
&unsqueeze252_out1[..],
&constant1967_out1[..],
&constant1968_out1[..],
]
.concat()
.try_into()
.unwrap();
let reshape149_out1 = transpose120_out1.reshape(concat141_out1);
let linear107_out1 = self.linear107.forward(reshape149_out1);
let add92_out1 = add88_out1.add(linear107_out1);
add92_out1
}
}
#[derive(Module, Debug)]
pub struct Submodule17 {
layernormalization58: LayerNorm,
conv1d34: Conv1d,
conv1d35: Conv1d,
conv1d36: Conv1d,
layernormalization59: LayerNorm,
linear108: Linear,
linear109: Linear,
constant1978: burn::module::Param<Tensor<1>>,
layernormalization60: LayerNorm,
layernormalization61: LayerNorm,
linear110: Linear,
linear111: Linear,
constant1979: burn::module::Param<Tensor<1>>,
layernormalization62: LayerNorm,
linear112: Linear,
linear113: Linear,
linear114: Linear,
linear115: Linear,
constant312: burn::module::Param<Tensor<2>>,
constant313: burn::module::Param<Tensor<2>>,
constant2025: burn::module::Param<Tensor<1>>,
linear116: Linear,
#[module(skip)]
device: Device,
}
impl Submodule17 {
#[allow(unused_variables)]
pub fn new(device: &Device) -> Self {
let layernormalization58 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let conv1d34 = Conv1dConfig::new(1024, 2048, 1)
.with_stride(1)
.with_padding(PaddingConfig1d::Valid)
.with_dilation(1)
.with_groups(1)
.with_bias(true)
.init(device);
let conv1d35 = Conv1dConfig::new(1024, 1024, 9)
.with_stride(1)
.with_padding(PaddingConfig1d::Valid)
.with_dilation(1)
.with_groups(1024)
.with_bias(true)
.init(device);
let conv1d36 = Conv1dConfig::new(1024, 1024, 1)
.with_stride(1)
.with_padding(PaddingConfig1d::Valid)
.with_dilation(1)
.with_groups(1)
.with_bias(true)
.init(device);
let layernormalization59 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let linear108 = LinearConfig::new(1024, 4096).with_bias(true).init(device);
let linear109 = LinearConfig::new(4096, 1024).with_bias(true).init(device);
let constant1978: 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([0.5f64]),
(device, burn::tensor::DType::F32),
),
device.clone(),
false,
[1].into(),
);
let layernormalization60 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let layernormalization61 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let linear110 = LinearConfig::new(1024, 4096).with_bias(true).init(device);
let linear111 = LinearConfig::new(4096, 1024).with_bias(true).init(device);
let constant1979: 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([0.5f64]),
(device, burn::tensor::DType::F32),
),
device.clone(),
false,
[1].into(),
);
let layernormalization62 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let linear112 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
let linear113 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
let linear114 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
let linear115 = LinearConfig::new(1024, 1024).with_bias(false).init(device);
let constant312: burn::module::Param<Tensor<2>> = burn::module::Param::uninitialized(
burn::module::ParamId::new(),
move |device, _require_grad| Tensor::<
2,
>::zeros([8, 128], (device, burn::tensor::DType::F32)),
device.clone(),
false,
[8, 128].into(),
);
let constant313: burn::module::Param<Tensor<2>> = burn::module::Param::uninitialized(
burn::module::ParamId::new(),
move |device, _require_grad| Tensor::<
2,
>::zeros([8, 128], (device, burn::tensor::DType::F32)),
device.clone(),
false,
[8, 128].into(),
);
let constant2025: 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([11.313708305358887f64]),
(device, burn::tensor::DType::F32),
),
device.clone(),
false,
[1].into(),
);
let linear116 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
Self {
layernormalization58,
conv1d34,
conv1d35,
conv1d36,
layernormalization59,
linear108,
linear109,
constant1978,
layernormalization60,
layernormalization61,
linear110,
linear111,
constant1979,
layernormalization62,
linear112,
linear113,
linear114,
linear115,
constant312,
constant313,
constant2025,
linear116,
device: device.clone(),
}
}
#[allow(clippy::let_and_return, clippy::approx_constant)]
pub fn forward(
&self,
add92_out1: Tensor<3>,
unsqueeze120_out1: Tensor<3, Bool>,
slice1_out1: Tensor<3>,
unsqueeze118_out1: Tensor<4, Bool>,
) -> Tensor<3> {
let layernormalization58_out1 = {
let dtype = add92_out1.clone().dtype();
self.layernormalization58
.forward(add92_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let transpose121_out1 = layernormalization58_out1.permute([0, 2, 1]);
let conv1d34_out1 = self.conv1d34.forward(transpose121_out1);
let split_tensors = conv1d34_out1.split_with_sizes([1024, 1024].into(), 1);
let [split12_out1, split12_out2] = split_tensors.try_into().unwrap();
let sigmoid46_out1 = burn::tensor::activation::sigmoid(split12_out2);
let mul115_out1 = split12_out1.mul(sigmoid46_out1);
let constant1969_out1 = 0f32;
let where54_out1 = mul115_out1.mask_fill(unsqueeze120_out1, constant1969_out1);
let pad24_out1 = where54_out1
.pad(
[(0usize, 0usize), (0usize, 0usize), (4usize, 4usize)],
burn::tensor::ops::PadMode::Constant(0f32),
);
let conv1d35_out1 = self.conv1d35.forward(pad24_out1);
let sigmoid47_out1 = burn::tensor::activation::sigmoid(conv1d35_out1.clone());
let mul116_out1 = conv1d35_out1.mul(sigmoid47_out1);
let conv1d36_out1 = self.conv1d36.forward(mul116_out1);
let transpose123_out1 = conv1d36_out1.permute([0, 2, 1]);
let add93_out1 = add92_out1.add(transpose123_out1);
let layernormalization59_out1 = {
let dtype = add93_out1.clone().dtype();
self.layernormalization59
.forward(add93_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let linear108_out1 = self.linear108.forward(layernormalization59_out1);
let sigmoid48_out1 = burn::tensor::activation::sigmoid(linear108_out1.clone());
let mul117_out1 = linear108_out1.mul(sigmoid48_out1);
let linear109_out1 = self.linear109.forward(mul117_out1);
let constant1978_out1 = self.constant1978.val();
let mul118_out1 = linear109_out1
.mul((constant1978_out1).unsqueeze_dims(&[0isize, 1isize]));
let add94_out1 = add93_out1.add(mul118_out1);
let layernormalization60_out1 = {
let dtype = add94_out1.dtype();
self.layernormalization60
.forward(add94_out1.cast(burn::tensor::DType::F32))
.cast(dtype)
};
let layernormalization61_out1 = {
let dtype = layernormalization60_out1.clone().dtype();
self.layernormalization61
.forward(
layernormalization60_out1.clone().cast(burn::tensor::DType::F32),
)
.cast(dtype)
};
let linear110_out1 = self.linear110.forward(layernormalization61_out1);
let sigmoid49_out1 = burn::tensor::activation::sigmoid(linear110_out1.clone());
let mul119_out1 = linear110_out1.mul(sigmoid49_out1);
let linear111_out1 = self.linear111.forward(mul119_out1);
let constant1979_out1 = self.constant1979.val();
let mul120_out1 = linear111_out1
.mul((constant1979_out1).unsqueeze_dims(&[0isize, 1isize]));
let add95_out1 = layernormalization60_out1.add(mul120_out1);
let layernormalization62_out1 = {
let dtype = add95_out1.clone().dtype();
self.layernormalization62
.forward(add95_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let shape146_out1: [i64; 3] = {
let axes = &layernormalization62_out1.clone().dims()[0..3];
let mut output = [0i64; 3];
for i in 0..3 {
output[i] = axes[i] as i64;
}
output
};
let gather115_out1 = shape146_out1[0] as i64;
let linear112_out1 = self.linear112.forward(layernormalization62_out1.clone());
let unsqueeze253_out1 = [gather115_out1 as i64];
let constant1984_out1: [i64; 1] = [128i64];
let constant1982_out1: [i64; 1] = [-1i64];
let constant1983_out1: [i64; 1] = [8i64];
let concat143_out1: [i64; 4usize] = [
&unsqueeze253_out1[..],
&constant1982_out1[..],
&constant1983_out1[..],
&constant1984_out1[..],
]
.concat()
.try_into()
.unwrap();
let unsqueeze254_out1 = [gather115_out1 as i64];
let constant1988_out1: [i64; 1] = [128i64];
let constant1986_out1: [i64; 1] = [-1i64];
let constant1987_out1: [i64; 1] = [8i64];
let concat144_out1: [i64; 4usize] = [
&unsqueeze254_out1[..],
&constant1986_out1[..],
&constant1987_out1[..],
&constant1988_out1[..],
]
.concat()
.try_into()
.unwrap();
let unsqueeze255_out1 = [gather115_out1 as i64];
let constant1992_out1: [i64; 1] = [128i64];
let constant1990_out1: [i64; 1] = [-1i64];
let constant1991_out1: [i64; 1] = [8i64];
let concat145_out1: [i64; 4usize] = [
&unsqueeze255_out1[..],
&constant1990_out1[..],
&constant1991_out1[..],
&constant1992_out1[..],
]
.concat()
.try_into()
.unwrap();
let reshape152_out1 = linear112_out1.reshape(concat143_out1);
let linear113_out1 = self.linear113.forward(layernormalization62_out1.clone());
let reshape153_out1 = linear113_out1.reshape(concat144_out1);
let linear114_out1 = self.linear114.forward(layernormalization62_out1);
let reshape154_out1 = linear114_out1.reshape(concat145_out1);
let transpose124_out1 = reshape154_out1.permute([0, 2, 1, 3]);
let linear115_out1 = self.linear115.forward(slice1_out1);
let reshape155_out1 = linear115_out1.reshape([1, -1, 8, 128]);
let constant312_out1 = self.constant312.val();
let add96_out1 = reshape152_out1
.clone()
.add((constant312_out1).unsqueeze_dims(&[0isize, 1isize]));
let transpose125_out1 = add96_out1.permute([0, 2, 1, 3]);
let constant313_out1 = self.constant313.val();
let add97_out1 = reshape152_out1
.add((constant313_out1).unsqueeze_dims(&[0isize, 1isize]));
let transpose126_out1 = add97_out1.permute([0, 2, 1, 3]);
let transpose127_out1 = reshape155_out1.permute([0, 2, 3, 1]);
let matmul152_out1 = transpose126_out1.matmul(transpose127_out1);
let shape147_out1: [i64; 4] = {
let axes = &matmul152_out1.clone().dims()[0..4];
let mut output = [0i64; 4];
for i in 0..4 {
output[i] = axes[i] as i64;
}
output
};
let gather116_out1 = shape147_out1[0] as i64;
let gather117_out1 = shape147_out1[1] as i64;
let gather118_out1 = shape147_out1[2] as i64;
let gather119_out1 = shape147_out1[3] as i64;
let pad25_out1 = matmul152_out1
.pad(
[(0usize, 0usize), (0usize, 0usize), (0usize, 0usize), (1usize, 0usize)],
burn::tensor::ops::PadMode::Constant(0f32),
);
let unsqueeze256_out1 = [gather116_out1 as i64];
let unsqueeze257_out1 = [gather117_out1 as i64];
let unsqueeze258_out1 = [gather118_out1 as i64];
let constant2007_out1: [i64; 1] = [-1i64];
let concat147_out1: [i64; 4usize] = [
&unsqueeze256_out1[..],
&unsqueeze257_out1[..],
&constant2007_out1[..],
&unsqueeze258_out1[..],
]
.concat()
.try_into()
.unwrap();
let reshape158_out1 = pad25_out1.reshape(concat147_out1);
let slice65_out1 = reshape158_out1.slice(s![.., .., 1.., ..]);
let unsqueeze259_out1 = [gather116_out1 as i64];
let unsqueeze260_out1 = [gather117_out1 as i64];
let unsqueeze261_out1 = [gather118_out1 as i64];
let unsqueeze262_out1 = [gather119_out1 as i64];
let concat148_out1: [i64; 4usize] = [
&unsqueeze259_out1[..],
&unsqueeze260_out1[..],
&unsqueeze261_out1[..],
&unsqueeze262_out1[..],
]
.concat()
.try_into()
.unwrap();
let reshape159_out1 = slice65_out1.reshape(concat148_out1);
let transpose129_out1 = reshape153_out1.permute([0, 2, 3, 1]);
let matmul153_out1 = transpose125_out1.matmul(transpose129_out1);
let shape151_out1: [i64; 4] = {
let axes = &matmul153_out1.clone().dims()[0..4];
let mut output = [0i64; 4];
for i in 0..4 {
output[i] = axes[i] as i64;
}
output
};
let slice66_out1: [i64; 1] = shape151_out1[3..4].try_into().unwrap();
let squeeze13_out1 = slice66_out1[0] as i64;
let unsqueeze263_out1 = [squeeze13_out1 as i64];
let slice67_out1 = reshape159_out1
.slice(s![.., .., .., 0..unsqueeze263_out1[0]]);
let add98_out1 = matmul153_out1.add(slice67_out1);
let constant2025_out1 = self.constant2025.val();
let div16_out1 = add98_out1
.div((constant2025_out1).unsqueeze_dims(&[0isize, 1isize, 2isize]));
let shape152_out1: [i64; 4] = {
let axes = &transpose124_out1.clone().dims()[0..4];
let mut output = [0i64; 4];
for i in 0..4 {
output[i] = axes[i] as i64;
}
output
};
let gather120_out1 = shape152_out1[0] as i64;
let constant2027_out1 = -10000f32;
let where55_out1 = div16_out1
.mask_fill(unsqueeze118_out1.clone(), constant2027_out1);
let softmax13_out1 = burn::tensor::activation::softmax(where55_out1, 3);
let constant2028_out1 = 0f32;
let where56_out1 = softmax13_out1
.mask_fill(unsqueeze118_out1, constant2028_out1);
let matmul154_out1 = where56_out1.matmul(transpose124_out1);
let transpose130_out1 = matmul154_out1.permute([0, 2, 1, 3]);
let unsqueeze264_out1 = [gather120_out1 as i64];
let constant2031_out1: [i64; 1] = [1024i64];
let constant2030_out1: [i64; 1] = [-1i64];
let concat149_out1: [i64; 3usize] = [
&unsqueeze264_out1[..],
&constant2030_out1[..],
&constant2031_out1[..],
]
.concat()
.try_into()
.unwrap();
let reshape160_out1 = transpose130_out1.reshape(concat149_out1);
let linear116_out1 = self.linear116.forward(reshape160_out1);
let add99_out1 = add95_out1.add(linear116_out1);
add99_out1
}
}
#[derive(Module, Debug)]
pub struct Submodule18 {
layernormalization63: LayerNorm,
conv1d37: Conv1d,
conv1d38: Conv1d,
conv1d39: Conv1d,
layernormalization64: LayerNorm,
linear117: Linear,
linear118: Linear,
constant2041: burn::module::Param<Tensor<1>>,
layernormalization65: LayerNorm,
layernormalization66: LayerNorm,
linear119: Linear,
linear120: Linear,
constant2042: burn::module::Param<Tensor<1>>,
layernormalization67: LayerNorm,
linear121: Linear,
linear122: Linear,
linear123: Linear,
linear124: Linear,
constant336: burn::module::Param<Tensor<2>>,
constant337: burn::module::Param<Tensor<2>>,
constant2088: burn::module::Param<Tensor<1>>,
linear125: Linear,
#[module(skip)]
device: Device,
}
impl Submodule18 {
#[allow(unused_variables)]
pub fn new(device: &Device) -> Self {
let layernormalization63 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let conv1d37 = Conv1dConfig::new(1024, 2048, 1)
.with_stride(1)
.with_padding(PaddingConfig1d::Valid)
.with_dilation(1)
.with_groups(1)
.with_bias(true)
.init(device);
let conv1d38 = Conv1dConfig::new(1024, 1024, 9)
.with_stride(1)
.with_padding(PaddingConfig1d::Valid)
.with_dilation(1)
.with_groups(1024)
.with_bias(true)
.init(device);
let conv1d39 = Conv1dConfig::new(1024, 1024, 1)
.with_stride(1)
.with_padding(PaddingConfig1d::Valid)
.with_dilation(1)
.with_groups(1)
.with_bias(true)
.init(device);
let layernormalization64 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let linear117 = LinearConfig::new(1024, 4096).with_bias(true).init(device);
let linear118 = LinearConfig::new(4096, 1024).with_bias(true).init(device);
let constant2041: 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([0.5f64]),
(device, burn::tensor::DType::F32),
),
device.clone(),
false,
[1].into(),
);
let layernormalization65 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let layernormalization66 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let linear119 = LinearConfig::new(1024, 4096).with_bias(true).init(device);
let linear120 = LinearConfig::new(4096, 1024).with_bias(true).init(device);
let constant2042: 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([0.5f64]),
(device, burn::tensor::DType::F32),
),
device.clone(),
false,
[1].into(),
);
let layernormalization67 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let linear121 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
let linear122 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
let linear123 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
let linear124 = LinearConfig::new(1024, 1024).with_bias(false).init(device);
let constant336: burn::module::Param<Tensor<2>> = burn::module::Param::uninitialized(
burn::module::ParamId::new(),
move |device, _require_grad| Tensor::<
2,
>::zeros([8, 128], (device, burn::tensor::DType::F32)),
device.clone(),
false,
[8, 128].into(),
);
let constant337: burn::module::Param<Tensor<2>> = burn::module::Param::uninitialized(
burn::module::ParamId::new(),
move |device, _require_grad| Tensor::<
2,
>::zeros([8, 128], (device, burn::tensor::DType::F32)),
device.clone(),
false,
[8, 128].into(),
);
let constant2088: 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([11.313708305358887f64]),
(device, burn::tensor::DType::F32),
),
device.clone(),
false,
[1].into(),
);
let linear125 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
Self {
layernormalization63,
conv1d37,
conv1d38,
conv1d39,
layernormalization64,
linear117,
linear118,
constant2041,
layernormalization65,
layernormalization66,
linear119,
linear120,
constant2042,
layernormalization67,
linear121,
linear122,
linear123,
linear124,
constant336,
constant337,
constant2088,
linear125,
device: device.clone(),
}
}
#[allow(clippy::let_and_return, clippy::approx_constant)]
pub fn forward(
&self,
add99_out1: Tensor<3>,
unsqueeze120_out1: Tensor<3, Bool>,
slice1_out1: Tensor<3>,
unsqueeze118_out1: Tensor<4, Bool>,
) -> Tensor<3> {
let layernormalization63_out1 = {
let dtype = add99_out1.clone().dtype();
self.layernormalization63
.forward(add99_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let transpose131_out1 = layernormalization63_out1.permute([0, 2, 1]);
let conv1d37_out1 = self.conv1d37.forward(transpose131_out1);
let split_tensors = conv1d37_out1.split_with_sizes([1024, 1024].into(), 1);
let [split13_out1, split13_out2] = split_tensors.try_into().unwrap();
let sigmoid50_out1 = burn::tensor::activation::sigmoid(split13_out2);
let mul121_out1 = split13_out1.mul(sigmoid50_out1);
let constant2032_out1 = 0f32;
let where57_out1 = mul121_out1.mask_fill(unsqueeze120_out1, constant2032_out1);
let pad26_out1 = where57_out1
.pad(
[(0usize, 0usize), (0usize, 0usize), (4usize, 4usize)],
burn::tensor::ops::PadMode::Constant(0f32),
);
let conv1d38_out1 = self.conv1d38.forward(pad26_out1);
let sigmoid51_out1 = burn::tensor::activation::sigmoid(conv1d38_out1.clone());
let mul122_out1 = conv1d38_out1.mul(sigmoid51_out1);
let conv1d39_out1 = self.conv1d39.forward(mul122_out1);
let transpose133_out1 = conv1d39_out1.permute([0, 2, 1]);
let add100_out1 = add99_out1.add(transpose133_out1);
let layernormalization64_out1 = {
let dtype = add100_out1.clone().dtype();
self.layernormalization64
.forward(add100_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let linear117_out1 = self.linear117.forward(layernormalization64_out1);
let sigmoid52_out1 = burn::tensor::activation::sigmoid(linear117_out1.clone());
let mul123_out1 = linear117_out1.mul(sigmoid52_out1);
let linear118_out1 = self.linear118.forward(mul123_out1);
let constant2041_out1 = self.constant2041.val();
let mul124_out1 = linear118_out1
.mul((constant2041_out1).unsqueeze_dims(&[0isize, 1isize]));
let add101_out1 = add100_out1.add(mul124_out1);
let layernormalization65_out1 = {
let dtype = add101_out1.dtype();
self.layernormalization65
.forward(add101_out1.cast(burn::tensor::DType::F32))
.cast(dtype)
};
let layernormalization66_out1 = {
let dtype = layernormalization65_out1.clone().dtype();
self.layernormalization66
.forward(
layernormalization65_out1.clone().cast(burn::tensor::DType::F32),
)
.cast(dtype)
};
let linear119_out1 = self.linear119.forward(layernormalization66_out1);
let sigmoid53_out1 = burn::tensor::activation::sigmoid(linear119_out1.clone());
let mul125_out1 = linear119_out1.mul(sigmoid53_out1);
let linear120_out1 = self.linear120.forward(mul125_out1);
let constant2042_out1 = self.constant2042.val();
let mul126_out1 = linear120_out1
.mul((constant2042_out1).unsqueeze_dims(&[0isize, 1isize]));
let add102_out1 = layernormalization65_out1.add(mul126_out1);
let layernormalization67_out1 = {
let dtype = add102_out1.clone().dtype();
self.layernormalization67
.forward(add102_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let shape153_out1: [i64; 3] = {
let axes = &layernormalization67_out1.clone().dims()[0..3];
let mut output = [0i64; 3];
for i in 0..3 {
output[i] = axes[i] as i64;
}
output
};
let gather121_out1 = shape153_out1[0] as i64;
let linear121_out1 = self.linear121.forward(layernormalization67_out1.clone());
let unsqueeze265_out1 = [gather121_out1 as i64];
let constant2047_out1: [i64; 1] = [128i64];
let constant2045_out1: [i64; 1] = [-1i64];
let constant2046_out1: [i64; 1] = [8i64];
let concat151_out1: [i64; 4usize] = [
&unsqueeze265_out1[..],
&constant2045_out1[..],
&constant2046_out1[..],
&constant2047_out1[..],
]
.concat()
.try_into()
.unwrap();
let unsqueeze266_out1 = [gather121_out1 as i64];
let constant2051_out1: [i64; 1] = [128i64];
let constant2049_out1: [i64; 1] = [-1i64];
let constant2050_out1: [i64; 1] = [8i64];
let concat152_out1: [i64; 4usize] = [
&unsqueeze266_out1[..],
&constant2049_out1[..],
&constant2050_out1[..],
&constant2051_out1[..],
]
.concat()
.try_into()
.unwrap();
let unsqueeze267_out1 = [gather121_out1 as i64];
let constant2055_out1: [i64; 1] = [128i64];
let constant2053_out1: [i64; 1] = [-1i64];
let constant2054_out1: [i64; 1] = [8i64];
let concat153_out1: [i64; 4usize] = [
&unsqueeze267_out1[..],
&constant2053_out1[..],
&constant2054_out1[..],
&constant2055_out1[..],
]
.concat()
.try_into()
.unwrap();
let reshape163_out1 = linear121_out1.reshape(concat151_out1);
let linear122_out1 = self.linear122.forward(layernormalization67_out1.clone());
let reshape164_out1 = linear122_out1.reshape(concat152_out1);
let linear123_out1 = self.linear123.forward(layernormalization67_out1);
let reshape165_out1 = linear123_out1.reshape(concat153_out1);
let transpose134_out1 = reshape165_out1.permute([0, 2, 1, 3]);
let linear124_out1 = self.linear124.forward(slice1_out1);
let reshape166_out1 = linear124_out1.reshape([1, -1, 8, 128]);
let constant336_out1 = self.constant336.val();
let add103_out1 = reshape163_out1
.clone()
.add((constant336_out1).unsqueeze_dims(&[0isize, 1isize]));
let transpose135_out1 = add103_out1.permute([0, 2, 1, 3]);
let constant337_out1 = self.constant337.val();
let add104_out1 = reshape163_out1
.add((constant337_out1).unsqueeze_dims(&[0isize, 1isize]));
let transpose136_out1 = add104_out1.permute([0, 2, 1, 3]);
let transpose137_out1 = reshape166_out1.permute([0, 2, 3, 1]);
let matmul164_out1 = transpose136_out1.matmul(transpose137_out1);
let shape154_out1: [i64; 4] = {
let axes = &matmul164_out1.clone().dims()[0..4];
let mut output = [0i64; 4];
for i in 0..4 {
output[i] = axes[i] as i64;
}
output
};
let gather122_out1 = shape154_out1[0] as i64;
let gather123_out1 = shape154_out1[1] as i64;
let gather124_out1 = shape154_out1[2] as i64;
let gather125_out1 = shape154_out1[3] as i64;
let pad27_out1 = matmul164_out1
.pad(
[(0usize, 0usize), (0usize, 0usize), (0usize, 0usize), (1usize, 0usize)],
burn::tensor::ops::PadMode::Constant(0f32),
);
let unsqueeze268_out1 = [gather122_out1 as i64];
let unsqueeze269_out1 = [gather123_out1 as i64];
let unsqueeze270_out1 = [gather124_out1 as i64];
let constant2070_out1: [i64; 1] = [-1i64];
let concat155_out1: [i64; 4usize] = [
&unsqueeze268_out1[..],
&unsqueeze269_out1[..],
&constant2070_out1[..],
&unsqueeze270_out1[..],
]
.concat()
.try_into()
.unwrap();
let reshape169_out1 = pad27_out1.reshape(concat155_out1);
let slice70_out1 = reshape169_out1.slice(s![.., .., 1.., ..]);
let unsqueeze271_out1 = [gather122_out1 as i64];
let unsqueeze272_out1 = [gather123_out1 as i64];
let unsqueeze273_out1 = [gather124_out1 as i64];
let unsqueeze274_out1 = [gather125_out1 as i64];
let concat156_out1: [i64; 4usize] = [
&unsqueeze271_out1[..],
&unsqueeze272_out1[..],
&unsqueeze273_out1[..],
&unsqueeze274_out1[..],
]
.concat()
.try_into()
.unwrap();
let reshape170_out1 = slice70_out1.reshape(concat156_out1);
let transpose139_out1 = reshape164_out1.permute([0, 2, 3, 1]);
let matmul165_out1 = transpose135_out1.matmul(transpose139_out1);
let shape158_out1: [i64; 4] = {
let axes = &matmul165_out1.clone().dims()[0..4];
let mut output = [0i64; 4];
for i in 0..4 {
output[i] = axes[i] as i64;
}
output
};
let slice71_out1: [i64; 1] = shape158_out1[3..4].try_into().unwrap();
let squeeze14_out1 = slice71_out1[0] as i64;
let unsqueeze275_out1 = [squeeze14_out1 as i64];
let slice72_out1 = reshape170_out1
.slice(s![.., .., .., 0..unsqueeze275_out1[0]]);
let add105_out1 = matmul165_out1.add(slice72_out1);
let constant2088_out1 = self.constant2088.val();
let div17_out1 = add105_out1
.div((constant2088_out1).unsqueeze_dims(&[0isize, 1isize, 2isize]));
let shape159_out1: [i64; 4] = {
let axes = &transpose134_out1.clone().dims()[0..4];
let mut output = [0i64; 4];
for i in 0..4 {
output[i] = axes[i] as i64;
}
output
};
let gather126_out1 = shape159_out1[0] as i64;
let constant2090_out1 = -10000f32;
let where58_out1 = div17_out1
.mask_fill(unsqueeze118_out1.clone(), constant2090_out1);
let softmax14_out1 = burn::tensor::activation::softmax(where58_out1, 3);
let constant2091_out1 = 0f32;
let where59_out1 = softmax14_out1
.mask_fill(unsqueeze118_out1, constant2091_out1);
let matmul166_out1 = where59_out1.matmul(transpose134_out1);
let transpose140_out1 = matmul166_out1.permute([0, 2, 1, 3]);
let unsqueeze276_out1 = [gather126_out1 as i64];
let constant2094_out1: [i64; 1] = [1024i64];
let constant2093_out1: [i64; 1] = [-1i64];
let concat157_out1: [i64; 3usize] = [
&unsqueeze276_out1[..],
&constant2093_out1[..],
&constant2094_out1[..],
]
.concat()
.try_into()
.unwrap();
let reshape171_out1 = transpose140_out1.reshape(concat157_out1);
let linear125_out1 = self.linear125.forward(reshape171_out1);
let add106_out1 = add102_out1.add(linear125_out1);
add106_out1
}
}
#[derive(Module, Debug)]
pub struct Submodule19 {
layernormalization68: LayerNorm,
conv1d40: Conv1d,
conv1d41: Conv1d,
conv1d42: Conv1d,
layernormalization69: LayerNorm,
linear126: Linear,
linear127: Linear,
constant2104: burn::module::Param<Tensor<1>>,
layernormalization70: LayerNorm,
layernormalization71: LayerNorm,
linear128: Linear,
linear129: Linear,
constant2105: burn::module::Param<Tensor<1>>,
layernormalization72: LayerNorm,
linear130: Linear,
linear131: Linear,
linear132: Linear,
linear133: Linear,
constant360: burn::module::Param<Tensor<2>>,
constant361: burn::module::Param<Tensor<2>>,
constant2151: burn::module::Param<Tensor<1>>,
linear134: Linear,
#[module(skip)]
device: Device,
}
impl Submodule19 {
#[allow(unused_variables)]
pub fn new(device: &Device) -> Self {
let layernormalization68 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let conv1d40 = Conv1dConfig::new(1024, 2048, 1)
.with_stride(1)
.with_padding(PaddingConfig1d::Valid)
.with_dilation(1)
.with_groups(1)
.with_bias(true)
.init(device);
let conv1d41 = Conv1dConfig::new(1024, 1024, 9)
.with_stride(1)
.with_padding(PaddingConfig1d::Valid)
.with_dilation(1)
.with_groups(1024)
.with_bias(true)
.init(device);
let conv1d42 = Conv1dConfig::new(1024, 1024, 1)
.with_stride(1)
.with_padding(PaddingConfig1d::Valid)
.with_dilation(1)
.with_groups(1)
.with_bias(true)
.init(device);
let layernormalization69 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let linear126 = LinearConfig::new(1024, 4096).with_bias(true).init(device);
let linear127 = LinearConfig::new(4096, 1024).with_bias(true).init(device);
let constant2104: 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([0.5f64]),
(device, burn::tensor::DType::F32),
),
device.clone(),
false,
[1].into(),
);
let layernormalization70 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let layernormalization71 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let linear128 = LinearConfig::new(1024, 4096).with_bias(true).init(device);
let linear129 = LinearConfig::new(4096, 1024).with_bias(true).init(device);
let constant2105: 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([0.5f64]),
(device, burn::tensor::DType::F32),
),
device.clone(),
false,
[1].into(),
);
let layernormalization72 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let linear130 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
let linear131 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
let linear132 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
let linear133 = LinearConfig::new(1024, 1024).with_bias(false).init(device);
let constant360: burn::module::Param<Tensor<2>> = burn::module::Param::uninitialized(
burn::module::ParamId::new(),
move |device, _require_grad| Tensor::<
2,
>::zeros([8, 128], (device, burn::tensor::DType::F32)),
device.clone(),
false,
[8, 128].into(),
);
let constant361: burn::module::Param<Tensor<2>> = burn::module::Param::uninitialized(
burn::module::ParamId::new(),
move |device, _require_grad| Tensor::<
2,
>::zeros([8, 128], (device, burn::tensor::DType::F32)),
device.clone(),
false,
[8, 128].into(),
);
let constant2151: 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([11.313708305358887f64]),
(device, burn::tensor::DType::F32),
),
device.clone(),
false,
[1].into(),
);
let linear134 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
Self {
layernormalization68,
conv1d40,
conv1d41,
conv1d42,
layernormalization69,
linear126,
linear127,
constant2104,
layernormalization70,
layernormalization71,
linear128,
linear129,
constant2105,
layernormalization72,
linear130,
linear131,
linear132,
linear133,
constant360,
constant361,
constant2151,
linear134,
device: device.clone(),
}
}
#[allow(clippy::let_and_return, clippy::approx_constant)]
pub fn forward(
&self,
add106_out1: Tensor<3>,
unsqueeze120_out1: Tensor<3, Bool>,
slice1_out1: Tensor<3>,
unsqueeze118_out1: Tensor<4, Bool>,
) -> Tensor<3> {
let layernormalization68_out1 = {
let dtype = add106_out1.clone().dtype();
self.layernormalization68
.forward(add106_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let transpose141_out1 = layernormalization68_out1.permute([0, 2, 1]);
let conv1d40_out1 = self.conv1d40.forward(transpose141_out1);
let split_tensors = conv1d40_out1.split_with_sizes([1024, 1024].into(), 1);
let [split14_out1, split14_out2] = split_tensors.try_into().unwrap();
let sigmoid54_out1 = burn::tensor::activation::sigmoid(split14_out2);
let mul127_out1 = split14_out1.mul(sigmoid54_out1);
let constant2095_out1 = 0f32;
let where60_out1 = mul127_out1.mask_fill(unsqueeze120_out1, constant2095_out1);
let pad28_out1 = where60_out1
.pad(
[(0usize, 0usize), (0usize, 0usize), (4usize, 4usize)],
burn::tensor::ops::PadMode::Constant(0f32),
);
let conv1d41_out1 = self.conv1d41.forward(pad28_out1);
let sigmoid55_out1 = burn::tensor::activation::sigmoid(conv1d41_out1.clone());
let mul128_out1 = conv1d41_out1.mul(sigmoid55_out1);
let conv1d42_out1 = self.conv1d42.forward(mul128_out1);
let transpose143_out1 = conv1d42_out1.permute([0, 2, 1]);
let add107_out1 = add106_out1.add(transpose143_out1);
let layernormalization69_out1 = {
let dtype = add107_out1.clone().dtype();
self.layernormalization69
.forward(add107_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let linear126_out1 = self.linear126.forward(layernormalization69_out1);
let sigmoid56_out1 = burn::tensor::activation::sigmoid(linear126_out1.clone());
let mul129_out1 = linear126_out1.mul(sigmoid56_out1);
let linear127_out1 = self.linear127.forward(mul129_out1);
let constant2104_out1 = self.constant2104.val();
let mul130_out1 = linear127_out1
.mul((constant2104_out1).unsqueeze_dims(&[0isize, 1isize]));
let add108_out1 = add107_out1.add(mul130_out1);
let layernormalization70_out1 = {
let dtype = add108_out1.dtype();
self.layernormalization70
.forward(add108_out1.cast(burn::tensor::DType::F32))
.cast(dtype)
};
let layernormalization71_out1 = {
let dtype = layernormalization70_out1.clone().dtype();
self.layernormalization71
.forward(
layernormalization70_out1.clone().cast(burn::tensor::DType::F32),
)
.cast(dtype)
};
let linear128_out1 = self.linear128.forward(layernormalization71_out1);
let sigmoid57_out1 = burn::tensor::activation::sigmoid(linear128_out1.clone());
let mul131_out1 = linear128_out1.mul(sigmoid57_out1);
let linear129_out1 = self.linear129.forward(mul131_out1);
let constant2105_out1 = self.constant2105.val();
let mul132_out1 = linear129_out1
.mul((constant2105_out1).unsqueeze_dims(&[0isize, 1isize]));
let add109_out1 = layernormalization70_out1.add(mul132_out1);
let layernormalization72_out1 = {
let dtype = add109_out1.clone().dtype();
self.layernormalization72
.forward(add109_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let shape160_out1: [i64; 3] = {
let axes = &layernormalization72_out1.clone().dims()[0..3];
let mut output = [0i64; 3];
for i in 0..3 {
output[i] = axes[i] as i64;
}
output
};
let gather127_out1 = shape160_out1[0] as i64;
let linear130_out1 = self.linear130.forward(layernormalization72_out1.clone());
let unsqueeze277_out1 = [gather127_out1 as i64];
let constant2110_out1: [i64; 1] = [128i64];
let constant2108_out1: [i64; 1] = [-1i64];
let constant2109_out1: [i64; 1] = [8i64];
let concat159_out1: [i64; 4usize] = [
&unsqueeze277_out1[..],
&constant2108_out1[..],
&constant2109_out1[..],
&constant2110_out1[..],
]
.concat()
.try_into()
.unwrap();
let unsqueeze278_out1 = [gather127_out1 as i64];
let constant2114_out1: [i64; 1] = [128i64];
let constant2112_out1: [i64; 1] = [-1i64];
let constant2113_out1: [i64; 1] = [8i64];
let concat160_out1: [i64; 4usize] = [
&unsqueeze278_out1[..],
&constant2112_out1[..],
&constant2113_out1[..],
&constant2114_out1[..],
]
.concat()
.try_into()
.unwrap();
let unsqueeze279_out1 = [gather127_out1 as i64];
let constant2118_out1: [i64; 1] = [128i64];
let constant2116_out1: [i64; 1] = [-1i64];
let constant2117_out1: [i64; 1] = [8i64];
let concat161_out1: [i64; 4usize] = [
&unsqueeze279_out1[..],
&constant2116_out1[..],
&constant2117_out1[..],
&constant2118_out1[..],
]
.concat()
.try_into()
.unwrap();
let reshape174_out1 = linear130_out1.reshape(concat159_out1);
let linear131_out1 = self.linear131.forward(layernormalization72_out1.clone());
let reshape175_out1 = linear131_out1.reshape(concat160_out1);
let linear132_out1 = self.linear132.forward(layernormalization72_out1);
let reshape176_out1 = linear132_out1.reshape(concat161_out1);
let transpose144_out1 = reshape176_out1.permute([0, 2, 1, 3]);
let linear133_out1 = self.linear133.forward(slice1_out1);
let reshape177_out1 = linear133_out1.reshape([1, -1, 8, 128]);
let constant360_out1 = self.constant360.val();
let add110_out1 = reshape174_out1
.clone()
.add((constant360_out1).unsqueeze_dims(&[0isize, 1isize]));
let transpose145_out1 = add110_out1.permute([0, 2, 1, 3]);
let constant361_out1 = self.constant361.val();
let add111_out1 = reshape174_out1
.add((constant361_out1).unsqueeze_dims(&[0isize, 1isize]));
let transpose146_out1 = add111_out1.permute([0, 2, 1, 3]);
let transpose147_out1 = reshape177_out1.permute([0, 2, 3, 1]);
let matmul176_out1 = transpose146_out1.matmul(transpose147_out1);
let shape161_out1: [i64; 4] = {
let axes = &matmul176_out1.clone().dims()[0..4];
let mut output = [0i64; 4];
for i in 0..4 {
output[i] = axes[i] as i64;
}
output
};
let gather128_out1 = shape161_out1[0] as i64;
let gather129_out1 = shape161_out1[1] as i64;
let gather130_out1 = shape161_out1[2] as i64;
let gather131_out1 = shape161_out1[3] as i64;
let pad29_out1 = matmul176_out1
.pad(
[(0usize, 0usize), (0usize, 0usize), (0usize, 0usize), (1usize, 0usize)],
burn::tensor::ops::PadMode::Constant(0f32),
);
let unsqueeze280_out1 = [gather128_out1 as i64];
let unsqueeze281_out1 = [gather129_out1 as i64];
let unsqueeze282_out1 = [gather130_out1 as i64];
let constant2133_out1: [i64; 1] = [-1i64];
let concat163_out1: [i64; 4usize] = [
&unsqueeze280_out1[..],
&unsqueeze281_out1[..],
&constant2133_out1[..],
&unsqueeze282_out1[..],
]
.concat()
.try_into()
.unwrap();
let reshape180_out1 = pad29_out1.reshape(concat163_out1);
let slice75_out1 = reshape180_out1.slice(s![.., .., 1.., ..]);
let unsqueeze283_out1 = [gather128_out1 as i64];
let unsqueeze284_out1 = [gather129_out1 as i64];
let unsqueeze285_out1 = [gather130_out1 as i64];
let unsqueeze286_out1 = [gather131_out1 as i64];
let concat164_out1: [i64; 4usize] = [
&unsqueeze283_out1[..],
&unsqueeze284_out1[..],
&unsqueeze285_out1[..],
&unsqueeze286_out1[..],
]
.concat()
.try_into()
.unwrap();
let reshape181_out1 = slice75_out1.reshape(concat164_out1);
let transpose149_out1 = reshape175_out1.permute([0, 2, 3, 1]);
let matmul177_out1 = transpose145_out1.matmul(transpose149_out1);
let shape165_out1: [i64; 4] = {
let axes = &matmul177_out1.clone().dims()[0..4];
let mut output = [0i64; 4];
for i in 0..4 {
output[i] = axes[i] as i64;
}
output
};
let slice76_out1: [i64; 1] = shape165_out1[3..4].try_into().unwrap();
let squeeze15_out1 = slice76_out1[0] as i64;
let unsqueeze287_out1 = [squeeze15_out1 as i64];
let slice77_out1 = reshape181_out1
.slice(s![.., .., .., 0..unsqueeze287_out1[0]]);
let add112_out1 = matmul177_out1.add(slice77_out1);
let constant2151_out1 = self.constant2151.val();
let div18_out1 = add112_out1
.div((constant2151_out1).unsqueeze_dims(&[0isize, 1isize, 2isize]));
let shape166_out1: [i64; 4] = {
let axes = &transpose144_out1.clone().dims()[0..4];
let mut output = [0i64; 4];
for i in 0..4 {
output[i] = axes[i] as i64;
}
output
};
let gather132_out1 = shape166_out1[0] as i64;
let constant2153_out1 = -10000f32;
let where61_out1 = div18_out1
.mask_fill(unsqueeze118_out1.clone(), constant2153_out1);
let softmax15_out1 = burn::tensor::activation::softmax(where61_out1, 3);
let constant2154_out1 = 0f32;
let where62_out1 = softmax15_out1
.mask_fill(unsqueeze118_out1, constant2154_out1);
let matmul178_out1 = where62_out1.matmul(transpose144_out1);
let transpose150_out1 = matmul178_out1.permute([0, 2, 1, 3]);
let unsqueeze288_out1 = [gather132_out1 as i64];
let constant2157_out1: [i64; 1] = [1024i64];
let constant2156_out1: [i64; 1] = [-1i64];
let concat165_out1: [i64; 3usize] = [
&unsqueeze288_out1[..],
&constant2156_out1[..],
&constant2157_out1[..],
]
.concat()
.try_into()
.unwrap();
let reshape182_out1 = transpose150_out1.reshape(concat165_out1);
let linear134_out1 = self.linear134.forward(reshape182_out1);
let add113_out1 = add109_out1.add(linear134_out1);
add113_out1
}
}
#[derive(Module, Debug)]
pub struct Submodule20 {
layernormalization73: LayerNorm,
conv1d43: Conv1d,
conv1d44: Conv1d,
conv1d45: Conv1d,
layernormalization74: LayerNorm,
linear135: Linear,
linear136: Linear,
constant2167: burn::module::Param<Tensor<1>>,
layernormalization75: LayerNorm,
layernormalization76: LayerNorm,
linear137: Linear,
linear138: Linear,
constant2168: burn::module::Param<Tensor<1>>,
layernormalization77: LayerNorm,
linear139: Linear,
linear140: Linear,
linear141: Linear,
linear142: Linear,
constant384: burn::module::Param<Tensor<2>>,
constant385: burn::module::Param<Tensor<2>>,
constant2214: burn::module::Param<Tensor<1>>,
linear143: Linear,
#[module(skip)]
device: Device,
}
impl Submodule20 {
#[allow(unused_variables)]
pub fn new(device: &Device) -> Self {
let layernormalization73 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let conv1d43 = Conv1dConfig::new(1024, 2048, 1)
.with_stride(1)
.with_padding(PaddingConfig1d::Valid)
.with_dilation(1)
.with_groups(1)
.with_bias(true)
.init(device);
let conv1d44 = Conv1dConfig::new(1024, 1024, 9)
.with_stride(1)
.with_padding(PaddingConfig1d::Valid)
.with_dilation(1)
.with_groups(1024)
.with_bias(true)
.init(device);
let conv1d45 = Conv1dConfig::new(1024, 1024, 1)
.with_stride(1)
.with_padding(PaddingConfig1d::Valid)
.with_dilation(1)
.with_groups(1)
.with_bias(true)
.init(device);
let layernormalization74 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let linear135 = LinearConfig::new(1024, 4096).with_bias(true).init(device);
let linear136 = LinearConfig::new(4096, 1024).with_bias(true).init(device);
let constant2167: 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([0.5f64]),
(device, burn::tensor::DType::F32),
),
device.clone(),
false,
[1].into(),
);
let layernormalization75 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let layernormalization76 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let linear137 = LinearConfig::new(1024, 4096).with_bias(true).init(device);
let linear138 = LinearConfig::new(4096, 1024).with_bias(true).init(device);
let constant2168: 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([0.5f64]),
(device, burn::tensor::DType::F32),
),
device.clone(),
false,
[1].into(),
);
let layernormalization77 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let linear139 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
let linear140 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
let linear141 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
let linear142 = LinearConfig::new(1024, 1024).with_bias(false).init(device);
let constant384: burn::module::Param<Tensor<2>> = burn::module::Param::uninitialized(
burn::module::ParamId::new(),
move |device, _require_grad| Tensor::<
2,
>::zeros([8, 128], (device, burn::tensor::DType::F32)),
device.clone(),
false,
[8, 128].into(),
);
let constant385: burn::module::Param<Tensor<2>> = burn::module::Param::uninitialized(
burn::module::ParamId::new(),
move |device, _require_grad| Tensor::<
2,
>::zeros([8, 128], (device, burn::tensor::DType::F32)),
device.clone(),
false,
[8, 128].into(),
);
let constant2214: 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([11.313708305358887f64]),
(device, burn::tensor::DType::F32),
),
device.clone(),
false,
[1].into(),
);
let linear143 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
Self {
layernormalization73,
conv1d43,
conv1d44,
conv1d45,
layernormalization74,
linear135,
linear136,
constant2167,
layernormalization75,
layernormalization76,
linear137,
linear138,
constant2168,
layernormalization77,
linear139,
linear140,
linear141,
linear142,
constant384,
constant385,
constant2214,
linear143,
device: device.clone(),
}
}
#[allow(clippy::let_and_return, clippy::approx_constant)]
pub fn forward(
&self,
add113_out1: Tensor<3>,
unsqueeze120_out1: Tensor<3, Bool>,
slice1_out1: Tensor<3>,
unsqueeze118_out1: Tensor<4, Bool>,
) -> Tensor<3> {
let layernormalization73_out1 = {
let dtype = add113_out1.clone().dtype();
self.layernormalization73
.forward(add113_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let transpose151_out1 = layernormalization73_out1.permute([0, 2, 1]);
let conv1d43_out1 = self.conv1d43.forward(transpose151_out1);
let split_tensors = conv1d43_out1.split_with_sizes([1024, 1024].into(), 1);
let [split15_out1, split15_out2] = split_tensors.try_into().unwrap();
let sigmoid58_out1 = burn::tensor::activation::sigmoid(split15_out2);
let mul133_out1 = split15_out1.mul(sigmoid58_out1);
let constant2158_out1 = 0f32;
let where63_out1 = mul133_out1.mask_fill(unsqueeze120_out1, constant2158_out1);
let pad30_out1 = where63_out1
.pad(
[(0usize, 0usize), (0usize, 0usize), (4usize, 4usize)],
burn::tensor::ops::PadMode::Constant(0f32),
);
let conv1d44_out1 = self.conv1d44.forward(pad30_out1);
let sigmoid59_out1 = burn::tensor::activation::sigmoid(conv1d44_out1.clone());
let mul134_out1 = conv1d44_out1.mul(sigmoid59_out1);
let conv1d45_out1 = self.conv1d45.forward(mul134_out1);
let transpose153_out1 = conv1d45_out1.permute([0, 2, 1]);
let add114_out1 = add113_out1.add(transpose153_out1);
let layernormalization74_out1 = {
let dtype = add114_out1.clone().dtype();
self.layernormalization74
.forward(add114_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let linear135_out1 = self.linear135.forward(layernormalization74_out1);
let sigmoid60_out1 = burn::tensor::activation::sigmoid(linear135_out1.clone());
let mul135_out1 = linear135_out1.mul(sigmoid60_out1);
let linear136_out1 = self.linear136.forward(mul135_out1);
let constant2167_out1 = self.constant2167.val();
let mul136_out1 = linear136_out1
.mul((constant2167_out1).unsqueeze_dims(&[0isize, 1isize]));
let add115_out1 = add114_out1.add(mul136_out1);
let layernormalization75_out1 = {
let dtype = add115_out1.dtype();
self.layernormalization75
.forward(add115_out1.cast(burn::tensor::DType::F32))
.cast(dtype)
};
let layernormalization76_out1 = {
let dtype = layernormalization75_out1.clone().dtype();
self.layernormalization76
.forward(
layernormalization75_out1.clone().cast(burn::tensor::DType::F32),
)
.cast(dtype)
};
let linear137_out1 = self.linear137.forward(layernormalization76_out1);
let sigmoid61_out1 = burn::tensor::activation::sigmoid(linear137_out1.clone());
let mul137_out1 = linear137_out1.mul(sigmoid61_out1);
let linear138_out1 = self.linear138.forward(mul137_out1);
let constant2168_out1 = self.constant2168.val();
let mul138_out1 = linear138_out1
.mul((constant2168_out1).unsqueeze_dims(&[0isize, 1isize]));
let add116_out1 = layernormalization75_out1.add(mul138_out1);
let layernormalization77_out1 = {
let dtype = add116_out1.clone().dtype();
self.layernormalization77
.forward(add116_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let shape167_out1: [i64; 3] = {
let axes = &layernormalization77_out1.clone().dims()[0..3];
let mut output = [0i64; 3];
for i in 0..3 {
output[i] = axes[i] as i64;
}
output
};
let gather133_out1 = shape167_out1[0] as i64;
let linear139_out1 = self.linear139.forward(layernormalization77_out1.clone());
let unsqueeze289_out1 = [gather133_out1 as i64];
let constant2173_out1: [i64; 1] = [128i64];
let constant2171_out1: [i64; 1] = [-1i64];
let constant2172_out1: [i64; 1] = [8i64];
let concat167_out1: [i64; 4usize] = [
&unsqueeze289_out1[..],
&constant2171_out1[..],
&constant2172_out1[..],
&constant2173_out1[..],
]
.concat()
.try_into()
.unwrap();
let unsqueeze290_out1 = [gather133_out1 as i64];
let constant2177_out1: [i64; 1] = [128i64];
let constant2175_out1: [i64; 1] = [-1i64];
let constant2176_out1: [i64; 1] = [8i64];
let concat168_out1: [i64; 4usize] = [
&unsqueeze290_out1[..],
&constant2175_out1[..],
&constant2176_out1[..],
&constant2177_out1[..],
]
.concat()
.try_into()
.unwrap();
let unsqueeze291_out1 = [gather133_out1 as i64];
let constant2181_out1: [i64; 1] = [128i64];
let constant2179_out1: [i64; 1] = [-1i64];
let constant2180_out1: [i64; 1] = [8i64];
let concat169_out1: [i64; 4usize] = [
&unsqueeze291_out1[..],
&constant2179_out1[..],
&constant2180_out1[..],
&constant2181_out1[..],
]
.concat()
.try_into()
.unwrap();
let reshape185_out1 = linear139_out1.reshape(concat167_out1);
let linear140_out1 = self.linear140.forward(layernormalization77_out1.clone());
let reshape186_out1 = linear140_out1.reshape(concat168_out1);
let linear141_out1 = self.linear141.forward(layernormalization77_out1);
let reshape187_out1 = linear141_out1.reshape(concat169_out1);
let transpose154_out1 = reshape187_out1.permute([0, 2, 1, 3]);
let linear142_out1 = self.linear142.forward(slice1_out1);
let reshape188_out1 = linear142_out1.reshape([1, -1, 8, 128]);
let constant384_out1 = self.constant384.val();
let add117_out1 = reshape185_out1
.clone()
.add((constant384_out1).unsqueeze_dims(&[0isize, 1isize]));
let transpose155_out1 = add117_out1.permute([0, 2, 1, 3]);
let constant385_out1 = self.constant385.val();
let add118_out1 = reshape185_out1
.add((constant385_out1).unsqueeze_dims(&[0isize, 1isize]));
let transpose156_out1 = add118_out1.permute([0, 2, 1, 3]);
let transpose157_out1 = reshape188_out1.permute([0, 2, 3, 1]);
let matmul188_out1 = transpose156_out1.matmul(transpose157_out1);
let shape168_out1: [i64; 4] = {
let axes = &matmul188_out1.clone().dims()[0..4];
let mut output = [0i64; 4];
for i in 0..4 {
output[i] = axes[i] as i64;
}
output
};
let gather134_out1 = shape168_out1[0] as i64;
let gather135_out1 = shape168_out1[1] as i64;
let gather136_out1 = shape168_out1[2] as i64;
let gather137_out1 = shape168_out1[3] as i64;
let pad31_out1 = matmul188_out1
.pad(
[(0usize, 0usize), (0usize, 0usize), (0usize, 0usize), (1usize, 0usize)],
burn::tensor::ops::PadMode::Constant(0f32),
);
let unsqueeze292_out1 = [gather134_out1 as i64];
let unsqueeze293_out1 = [gather135_out1 as i64];
let unsqueeze294_out1 = [gather136_out1 as i64];
let constant2196_out1: [i64; 1] = [-1i64];
let concat171_out1: [i64; 4usize] = [
&unsqueeze292_out1[..],
&unsqueeze293_out1[..],
&constant2196_out1[..],
&unsqueeze294_out1[..],
]
.concat()
.try_into()
.unwrap();
let reshape191_out1 = pad31_out1.reshape(concat171_out1);
let slice80_out1 = reshape191_out1.slice(s![.., .., 1.., ..]);
let unsqueeze295_out1 = [gather134_out1 as i64];
let unsqueeze296_out1 = [gather135_out1 as i64];
let unsqueeze297_out1 = [gather136_out1 as i64];
let unsqueeze298_out1 = [gather137_out1 as i64];
let concat172_out1: [i64; 4usize] = [
&unsqueeze295_out1[..],
&unsqueeze296_out1[..],
&unsqueeze297_out1[..],
&unsqueeze298_out1[..],
]
.concat()
.try_into()
.unwrap();
let reshape192_out1 = slice80_out1.reshape(concat172_out1);
let transpose159_out1 = reshape186_out1.permute([0, 2, 3, 1]);
let matmul189_out1 = transpose155_out1.matmul(transpose159_out1);
let shape172_out1: [i64; 4] = {
let axes = &matmul189_out1.clone().dims()[0..4];
let mut output = [0i64; 4];
for i in 0..4 {
output[i] = axes[i] as i64;
}
output
};
let slice81_out1: [i64; 1] = shape172_out1[3..4].try_into().unwrap();
let squeeze16_out1 = slice81_out1[0] as i64;
let unsqueeze299_out1 = [squeeze16_out1 as i64];
let slice82_out1 = reshape192_out1
.slice(s![.., .., .., 0..unsqueeze299_out1[0]]);
let add119_out1 = matmul189_out1.add(slice82_out1);
let constant2214_out1 = self.constant2214.val();
let div19_out1 = add119_out1
.div((constant2214_out1).unsqueeze_dims(&[0isize, 1isize, 2isize]));
let shape173_out1: [i64; 4] = {
let axes = &transpose154_out1.clone().dims()[0..4];
let mut output = [0i64; 4];
for i in 0..4 {
output[i] = axes[i] as i64;
}
output
};
let gather138_out1 = shape173_out1[0] as i64;
let constant2216_out1 = -10000f32;
let where64_out1 = div19_out1
.mask_fill(unsqueeze118_out1.clone(), constant2216_out1);
let softmax16_out1 = burn::tensor::activation::softmax(where64_out1, 3);
let constant2217_out1 = 0f32;
let where65_out1 = softmax16_out1
.mask_fill(unsqueeze118_out1, constant2217_out1);
let matmul190_out1 = where65_out1.matmul(transpose154_out1);
let transpose160_out1 = matmul190_out1.permute([0, 2, 1, 3]);
let unsqueeze300_out1 = [gather138_out1 as i64];
let constant2220_out1: [i64; 1] = [1024i64];
let constant2219_out1: [i64; 1] = [-1i64];
let concat173_out1: [i64; 3usize] = [
&unsqueeze300_out1[..],
&constant2219_out1[..],
&constant2220_out1[..],
]
.concat()
.try_into()
.unwrap();
let reshape193_out1 = transpose160_out1.reshape(concat173_out1);
let linear143_out1 = self.linear143.forward(reshape193_out1);
let add120_out1 = add116_out1.add(linear143_out1);
add120_out1
}
}
#[derive(Module, Debug)]
pub struct Submodule21 {
layernormalization78: LayerNorm,
conv1d46: Conv1d,
conv1d47: Conv1d,
conv1d48: Conv1d,
layernormalization79: LayerNorm,
linear144: Linear,
linear145: Linear,
constant2230: burn::module::Param<Tensor<1>>,
layernormalization80: LayerNorm,
layernormalization81: LayerNorm,
linear146: Linear,
linear147: Linear,
constant2231: burn::module::Param<Tensor<1>>,
layernormalization82: LayerNorm,
linear148: Linear,
linear149: Linear,
linear150: Linear,
linear151: Linear,
constant408: burn::module::Param<Tensor<2>>,
constant409: burn::module::Param<Tensor<2>>,
constant2277: burn::module::Param<Tensor<1>>,
linear152: Linear,
#[module(skip)]
device: Device,
}
impl Submodule21 {
#[allow(unused_variables)]
pub fn new(device: &Device) -> Self {
let layernormalization78 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let conv1d46 = Conv1dConfig::new(1024, 2048, 1)
.with_stride(1)
.with_padding(PaddingConfig1d::Valid)
.with_dilation(1)
.with_groups(1)
.with_bias(true)
.init(device);
let conv1d47 = Conv1dConfig::new(1024, 1024, 9)
.with_stride(1)
.with_padding(PaddingConfig1d::Valid)
.with_dilation(1)
.with_groups(1024)
.with_bias(true)
.init(device);
let conv1d48 = Conv1dConfig::new(1024, 1024, 1)
.with_stride(1)
.with_padding(PaddingConfig1d::Valid)
.with_dilation(1)
.with_groups(1)
.with_bias(true)
.init(device);
let layernormalization79 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let linear144 = LinearConfig::new(1024, 4096).with_bias(true).init(device);
let linear145 = LinearConfig::new(4096, 1024).with_bias(true).init(device);
let constant2230: 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([0.5f64]),
(device, burn::tensor::DType::F32),
),
device.clone(),
false,
[1].into(),
);
let layernormalization80 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let layernormalization81 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let linear146 = LinearConfig::new(1024, 4096).with_bias(true).init(device);
let linear147 = LinearConfig::new(4096, 1024).with_bias(true).init(device);
let constant2231: 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([0.5f64]),
(device, burn::tensor::DType::F32),
),
device.clone(),
false,
[1].into(),
);
let layernormalization82 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let linear148 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
let linear149 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
let linear150 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
let linear151 = LinearConfig::new(1024, 1024).with_bias(false).init(device);
let constant408: burn::module::Param<Tensor<2>> = burn::module::Param::uninitialized(
burn::module::ParamId::new(),
move |device, _require_grad| Tensor::<
2,
>::zeros([8, 128], (device, burn::tensor::DType::F32)),
device.clone(),
false,
[8, 128].into(),
);
let constant409: burn::module::Param<Tensor<2>> = burn::module::Param::uninitialized(
burn::module::ParamId::new(),
move |device, _require_grad| Tensor::<
2,
>::zeros([8, 128], (device, burn::tensor::DType::F32)),
device.clone(),
false,
[8, 128].into(),
);
let constant2277: 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([11.313708305358887f64]),
(device, burn::tensor::DType::F32),
),
device.clone(),
false,
[1].into(),
);
let linear152 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
Self {
layernormalization78,
conv1d46,
conv1d47,
conv1d48,
layernormalization79,
linear144,
linear145,
constant2230,
layernormalization80,
layernormalization81,
linear146,
linear147,
constant2231,
layernormalization82,
linear148,
linear149,
linear150,
linear151,
constant408,
constant409,
constant2277,
linear152,
device: device.clone(),
}
}
#[allow(clippy::let_and_return, clippy::approx_constant)]
pub fn forward(
&self,
add120_out1: Tensor<3>,
unsqueeze120_out1: Tensor<3, Bool>,
slice1_out1: Tensor<3>,
unsqueeze118_out1: Tensor<4, Bool>,
) -> Tensor<3> {
let layernormalization78_out1 = {
let dtype = add120_out1.clone().dtype();
self.layernormalization78
.forward(add120_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let transpose161_out1 = layernormalization78_out1.permute([0, 2, 1]);
let conv1d46_out1 = self.conv1d46.forward(transpose161_out1);
let split_tensors = conv1d46_out1.split_with_sizes([1024, 1024].into(), 1);
let [split16_out1, split16_out2] = split_tensors.try_into().unwrap();
let sigmoid62_out1 = burn::tensor::activation::sigmoid(split16_out2);
let mul139_out1 = split16_out1.mul(sigmoid62_out1);
let constant2221_out1 = 0f32;
let where66_out1 = mul139_out1.mask_fill(unsqueeze120_out1, constant2221_out1);
let pad32_out1 = where66_out1
.pad(
[(0usize, 0usize), (0usize, 0usize), (4usize, 4usize)],
burn::tensor::ops::PadMode::Constant(0f32),
);
let conv1d47_out1 = self.conv1d47.forward(pad32_out1);
let sigmoid63_out1 = burn::tensor::activation::sigmoid(conv1d47_out1.clone());
let mul140_out1 = conv1d47_out1.mul(sigmoid63_out1);
let conv1d48_out1 = self.conv1d48.forward(mul140_out1);
let transpose163_out1 = conv1d48_out1.permute([0, 2, 1]);
let add121_out1 = add120_out1.add(transpose163_out1);
let layernormalization79_out1 = {
let dtype = add121_out1.clone().dtype();
self.layernormalization79
.forward(add121_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let linear144_out1 = self.linear144.forward(layernormalization79_out1);
let sigmoid64_out1 = burn::tensor::activation::sigmoid(linear144_out1.clone());
let mul141_out1 = linear144_out1.mul(sigmoid64_out1);
let linear145_out1 = self.linear145.forward(mul141_out1);
let constant2230_out1 = self.constant2230.val();
let mul142_out1 = linear145_out1
.mul((constant2230_out1).unsqueeze_dims(&[0isize, 1isize]));
let add122_out1 = add121_out1.add(mul142_out1);
let layernormalization80_out1 = {
let dtype = add122_out1.dtype();
self.layernormalization80
.forward(add122_out1.cast(burn::tensor::DType::F32))
.cast(dtype)
};
let layernormalization81_out1 = {
let dtype = layernormalization80_out1.clone().dtype();
self.layernormalization81
.forward(
layernormalization80_out1.clone().cast(burn::tensor::DType::F32),
)
.cast(dtype)
};
let linear146_out1 = self.linear146.forward(layernormalization81_out1);
let sigmoid65_out1 = burn::tensor::activation::sigmoid(linear146_out1.clone());
let mul143_out1 = linear146_out1.mul(sigmoid65_out1);
let linear147_out1 = self.linear147.forward(mul143_out1);
let constant2231_out1 = self.constant2231.val();
let mul144_out1 = linear147_out1
.mul((constant2231_out1).unsqueeze_dims(&[0isize, 1isize]));
let add123_out1 = layernormalization80_out1.add(mul144_out1);
let layernormalization82_out1 = {
let dtype = add123_out1.clone().dtype();
self.layernormalization82
.forward(add123_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let shape174_out1: [i64; 3] = {
let axes = &layernormalization82_out1.clone().dims()[0..3];
let mut output = [0i64; 3];
for i in 0..3 {
output[i] = axes[i] as i64;
}
output
};
let gather139_out1 = shape174_out1[0] as i64;
let linear148_out1 = self.linear148.forward(layernormalization82_out1.clone());
let unsqueeze301_out1 = [gather139_out1 as i64];
let constant2236_out1: [i64; 1] = [128i64];
let constant2234_out1: [i64; 1] = [-1i64];
let constant2235_out1: [i64; 1] = [8i64];
let concat175_out1: [i64; 4usize] = [
&unsqueeze301_out1[..],
&constant2234_out1[..],
&constant2235_out1[..],
&constant2236_out1[..],
]
.concat()
.try_into()
.unwrap();
let unsqueeze302_out1 = [gather139_out1 as i64];
let constant2240_out1: [i64; 1] = [128i64];
let constant2238_out1: [i64; 1] = [-1i64];
let constant2239_out1: [i64; 1] = [8i64];
let concat176_out1: [i64; 4usize] = [
&unsqueeze302_out1[..],
&constant2238_out1[..],
&constant2239_out1[..],
&constant2240_out1[..],
]
.concat()
.try_into()
.unwrap();
let unsqueeze303_out1 = [gather139_out1 as i64];
let constant2244_out1: [i64; 1] = [128i64];
let constant2242_out1: [i64; 1] = [-1i64];
let constant2243_out1: [i64; 1] = [8i64];
let concat177_out1: [i64; 4usize] = [
&unsqueeze303_out1[..],
&constant2242_out1[..],
&constant2243_out1[..],
&constant2244_out1[..],
]
.concat()
.try_into()
.unwrap();
let reshape196_out1 = linear148_out1.reshape(concat175_out1);
let linear149_out1 = self.linear149.forward(layernormalization82_out1.clone());
let reshape197_out1 = linear149_out1.reshape(concat176_out1);
let linear150_out1 = self.linear150.forward(layernormalization82_out1);
let reshape198_out1 = linear150_out1.reshape(concat177_out1);
let transpose164_out1 = reshape198_out1.permute([0, 2, 1, 3]);
let linear151_out1 = self.linear151.forward(slice1_out1);
let reshape199_out1 = linear151_out1.reshape([1, -1, 8, 128]);
let constant408_out1 = self.constant408.val();
let add124_out1 = reshape196_out1
.clone()
.add((constant408_out1).unsqueeze_dims(&[0isize, 1isize]));
let transpose165_out1 = add124_out1.permute([0, 2, 1, 3]);
let constant409_out1 = self.constant409.val();
let add125_out1 = reshape196_out1
.add((constant409_out1).unsqueeze_dims(&[0isize, 1isize]));
let transpose166_out1 = add125_out1.permute([0, 2, 1, 3]);
let transpose167_out1 = reshape199_out1.permute([0, 2, 3, 1]);
let matmul200_out1 = transpose166_out1.matmul(transpose167_out1);
let shape175_out1: [i64; 4] = {
let axes = &matmul200_out1.clone().dims()[0..4];
let mut output = [0i64; 4];
for i in 0..4 {
output[i] = axes[i] as i64;
}
output
};
let gather140_out1 = shape175_out1[0] as i64;
let gather141_out1 = shape175_out1[1] as i64;
let gather142_out1 = shape175_out1[2] as i64;
let gather143_out1 = shape175_out1[3] as i64;
let pad33_out1 = matmul200_out1
.pad(
[(0usize, 0usize), (0usize, 0usize), (0usize, 0usize), (1usize, 0usize)],
burn::tensor::ops::PadMode::Constant(0f32),
);
let unsqueeze304_out1 = [gather140_out1 as i64];
let unsqueeze305_out1 = [gather141_out1 as i64];
let unsqueeze306_out1 = [gather142_out1 as i64];
let constant2259_out1: [i64; 1] = [-1i64];
let concat179_out1: [i64; 4usize] = [
&unsqueeze304_out1[..],
&unsqueeze305_out1[..],
&constant2259_out1[..],
&unsqueeze306_out1[..],
]
.concat()
.try_into()
.unwrap();
let reshape202_out1 = pad33_out1.reshape(concat179_out1);
let slice85_out1 = reshape202_out1.slice(s![.., .., 1.., ..]);
let unsqueeze307_out1 = [gather140_out1 as i64];
let unsqueeze308_out1 = [gather141_out1 as i64];
let unsqueeze309_out1 = [gather142_out1 as i64];
let unsqueeze310_out1 = [gather143_out1 as i64];
let concat180_out1: [i64; 4usize] = [
&unsqueeze307_out1[..],
&unsqueeze308_out1[..],
&unsqueeze309_out1[..],
&unsqueeze310_out1[..],
]
.concat()
.try_into()
.unwrap();
let reshape203_out1 = slice85_out1.reshape(concat180_out1);
let transpose169_out1 = reshape197_out1.permute([0, 2, 3, 1]);
let matmul201_out1 = transpose165_out1.matmul(transpose169_out1);
let shape179_out1: [i64; 4] = {
let axes = &matmul201_out1.clone().dims()[0..4];
let mut output = [0i64; 4];
for i in 0..4 {
output[i] = axes[i] as i64;
}
output
};
let slice86_out1: [i64; 1] = shape179_out1[3..4].try_into().unwrap();
let squeeze17_out1 = slice86_out1[0] as i64;
let unsqueeze311_out1 = [squeeze17_out1 as i64];
let slice87_out1 = reshape203_out1
.slice(s![.., .., .., 0..unsqueeze311_out1[0]]);
let add126_out1 = matmul201_out1.add(slice87_out1);
let constant2277_out1 = self.constant2277.val();
let div20_out1 = add126_out1
.div((constant2277_out1).unsqueeze_dims(&[0isize, 1isize, 2isize]));
let shape180_out1: [i64; 4] = {
let axes = &transpose164_out1.clone().dims()[0..4];
let mut output = [0i64; 4];
for i in 0..4 {
output[i] = axes[i] as i64;
}
output
};
let gather144_out1 = shape180_out1[0] as i64;
let constant2279_out1 = -10000f32;
let where67_out1 = div20_out1
.mask_fill(unsqueeze118_out1.clone(), constant2279_out1);
let softmax17_out1 = burn::tensor::activation::softmax(where67_out1, 3);
let constant2280_out1 = 0f32;
let where68_out1 = softmax17_out1
.mask_fill(unsqueeze118_out1, constant2280_out1);
let matmul202_out1 = where68_out1.matmul(transpose164_out1);
let transpose170_out1 = matmul202_out1.permute([0, 2, 1, 3]);
let unsqueeze312_out1 = [gather144_out1 as i64];
let constant2283_out1: [i64; 1] = [1024i64];
let constant2282_out1: [i64; 1] = [-1i64];
let concat181_out1: [i64; 3usize] = [
&unsqueeze312_out1[..],
&constant2282_out1[..],
&constant2283_out1[..],
]
.concat()
.try_into()
.unwrap();
let reshape204_out1 = transpose170_out1.reshape(concat181_out1);
let linear152_out1 = self.linear152.forward(reshape204_out1);
let add127_out1 = add123_out1.add(linear152_out1);
add127_out1
}
}
#[derive(Module, Debug)]
pub struct Submodule22 {
layernormalization83: LayerNorm,
conv1d49: Conv1d,
conv1d50: Conv1d,
conv1d51: Conv1d,
layernormalization84: LayerNorm,
linear153: Linear,
linear154: Linear,
constant2293: burn::module::Param<Tensor<1>>,
layernormalization85: LayerNorm,
layernormalization86: LayerNorm,
linear155: Linear,
linear156: Linear,
constant2294: burn::module::Param<Tensor<1>>,
layernormalization87: LayerNorm,
linear157: Linear,
linear158: Linear,
linear159: Linear,
linear160: Linear,
constant432: burn::module::Param<Tensor<2>>,
constant433: burn::module::Param<Tensor<2>>,
constant2340: burn::module::Param<Tensor<1>>,
linear161: Linear,
#[module(skip)]
device: Device,
}
impl Submodule22 {
#[allow(unused_variables)]
pub fn new(device: &Device) -> Self {
let layernormalization83 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let conv1d49 = Conv1dConfig::new(1024, 2048, 1)
.with_stride(1)
.with_padding(PaddingConfig1d::Valid)
.with_dilation(1)
.with_groups(1)
.with_bias(true)
.init(device);
let conv1d50 = Conv1dConfig::new(1024, 1024, 9)
.with_stride(1)
.with_padding(PaddingConfig1d::Valid)
.with_dilation(1)
.with_groups(1024)
.with_bias(true)
.init(device);
let conv1d51 = Conv1dConfig::new(1024, 1024, 1)
.with_stride(1)
.with_padding(PaddingConfig1d::Valid)
.with_dilation(1)
.with_groups(1)
.with_bias(true)
.init(device);
let layernormalization84 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let linear153 = LinearConfig::new(1024, 4096).with_bias(true).init(device);
let linear154 = LinearConfig::new(4096, 1024).with_bias(true).init(device);
let constant2293: 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([0.5f64]),
(device, burn::tensor::DType::F32),
),
device.clone(),
false,
[1].into(),
);
let layernormalization85 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let layernormalization86 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let linear155 = LinearConfig::new(1024, 4096).with_bias(true).init(device);
let linear156 = LinearConfig::new(4096, 1024).with_bias(true).init(device);
let constant2294: 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([0.5f64]),
(device, burn::tensor::DType::F32),
),
device.clone(),
false,
[1].into(),
);
let layernormalization87 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let linear157 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
let linear158 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
let linear159 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
let linear160 = LinearConfig::new(1024, 1024).with_bias(false).init(device);
let constant432: burn::module::Param<Tensor<2>> = burn::module::Param::uninitialized(
burn::module::ParamId::new(),
move |device, _require_grad| Tensor::<
2,
>::zeros([8, 128], (device, burn::tensor::DType::F32)),
device.clone(),
false,
[8, 128].into(),
);
let constant433: burn::module::Param<Tensor<2>> = burn::module::Param::uninitialized(
burn::module::ParamId::new(),
move |device, _require_grad| Tensor::<
2,
>::zeros([8, 128], (device, burn::tensor::DType::F32)),
device.clone(),
false,
[8, 128].into(),
);
let constant2340: 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([11.313708305358887f64]),
(device, burn::tensor::DType::F32),
),
device.clone(),
false,
[1].into(),
);
let linear161 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
Self {
layernormalization83,
conv1d49,
conv1d50,
conv1d51,
layernormalization84,
linear153,
linear154,
constant2293,
layernormalization85,
layernormalization86,
linear155,
linear156,
constant2294,
layernormalization87,
linear157,
linear158,
linear159,
linear160,
constant432,
constant433,
constant2340,
linear161,
device: device.clone(),
}
}
#[allow(clippy::let_and_return, clippy::approx_constant)]
pub fn forward(
&self,
add127_out1: Tensor<3>,
unsqueeze120_out1: Tensor<3, Bool>,
slice1_out1: Tensor<3>,
unsqueeze118_out1: Tensor<4, Bool>,
) -> Tensor<3> {
let layernormalization83_out1 = {
let dtype = add127_out1.clone().dtype();
self.layernormalization83
.forward(add127_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let transpose171_out1 = layernormalization83_out1.permute([0, 2, 1]);
let conv1d49_out1 = self.conv1d49.forward(transpose171_out1);
let split_tensors = conv1d49_out1.split_with_sizes([1024, 1024].into(), 1);
let [split17_out1, split17_out2] = split_tensors.try_into().unwrap();
let sigmoid66_out1 = burn::tensor::activation::sigmoid(split17_out2);
let mul145_out1 = split17_out1.mul(sigmoid66_out1);
let constant2284_out1 = 0f32;
let where69_out1 = mul145_out1.mask_fill(unsqueeze120_out1, constant2284_out1);
let pad34_out1 = where69_out1
.pad(
[(0usize, 0usize), (0usize, 0usize), (4usize, 4usize)],
burn::tensor::ops::PadMode::Constant(0f32),
);
let conv1d50_out1 = self.conv1d50.forward(pad34_out1);
let sigmoid67_out1 = burn::tensor::activation::sigmoid(conv1d50_out1.clone());
let mul146_out1 = conv1d50_out1.mul(sigmoid67_out1);
let conv1d51_out1 = self.conv1d51.forward(mul146_out1);
let transpose173_out1 = conv1d51_out1.permute([0, 2, 1]);
let add128_out1 = add127_out1.add(transpose173_out1);
let layernormalization84_out1 = {
let dtype = add128_out1.clone().dtype();
self.layernormalization84
.forward(add128_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let linear153_out1 = self.linear153.forward(layernormalization84_out1);
let sigmoid68_out1 = burn::tensor::activation::sigmoid(linear153_out1.clone());
let mul147_out1 = linear153_out1.mul(sigmoid68_out1);
let linear154_out1 = self.linear154.forward(mul147_out1);
let constant2293_out1 = self.constant2293.val();
let mul148_out1 = linear154_out1
.mul((constant2293_out1).unsqueeze_dims(&[0isize, 1isize]));
let add129_out1 = add128_out1.add(mul148_out1);
let layernormalization85_out1 = {
let dtype = add129_out1.dtype();
self.layernormalization85
.forward(add129_out1.cast(burn::tensor::DType::F32))
.cast(dtype)
};
let layernormalization86_out1 = {
let dtype = layernormalization85_out1.clone().dtype();
self.layernormalization86
.forward(
layernormalization85_out1.clone().cast(burn::tensor::DType::F32),
)
.cast(dtype)
};
let linear155_out1 = self.linear155.forward(layernormalization86_out1);
let sigmoid69_out1 = burn::tensor::activation::sigmoid(linear155_out1.clone());
let mul149_out1 = linear155_out1.mul(sigmoid69_out1);
let linear156_out1 = self.linear156.forward(mul149_out1);
let constant2294_out1 = self.constant2294.val();
let mul150_out1 = linear156_out1
.mul((constant2294_out1).unsqueeze_dims(&[0isize, 1isize]));
let add130_out1 = layernormalization85_out1.add(mul150_out1);
let layernormalization87_out1 = {
let dtype = add130_out1.clone().dtype();
self.layernormalization87
.forward(add130_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let shape181_out1: [i64; 3] = {
let axes = &layernormalization87_out1.clone().dims()[0..3];
let mut output = [0i64; 3];
for i in 0..3 {
output[i] = axes[i] as i64;
}
output
};
let gather145_out1 = shape181_out1[0] as i64;
let linear157_out1 = self.linear157.forward(layernormalization87_out1.clone());
let unsqueeze313_out1 = [gather145_out1 as i64];
let constant2299_out1: [i64; 1] = [128i64];
let constant2297_out1: [i64; 1] = [-1i64];
let constant2298_out1: [i64; 1] = [8i64];
let concat183_out1: [i64; 4usize] = [
&unsqueeze313_out1[..],
&constant2297_out1[..],
&constant2298_out1[..],
&constant2299_out1[..],
]
.concat()
.try_into()
.unwrap();
let unsqueeze314_out1 = [gather145_out1 as i64];
let constant2303_out1: [i64; 1] = [128i64];
let constant2301_out1: [i64; 1] = [-1i64];
let constant2302_out1: [i64; 1] = [8i64];
let concat184_out1: [i64; 4usize] = [
&unsqueeze314_out1[..],
&constant2301_out1[..],
&constant2302_out1[..],
&constant2303_out1[..],
]
.concat()
.try_into()
.unwrap();
let unsqueeze315_out1 = [gather145_out1 as i64];
let constant2307_out1: [i64; 1] = [128i64];
let constant2305_out1: [i64; 1] = [-1i64];
let constant2306_out1: [i64; 1] = [8i64];
let concat185_out1: [i64; 4usize] = [
&unsqueeze315_out1[..],
&constant2305_out1[..],
&constant2306_out1[..],
&constant2307_out1[..],
]
.concat()
.try_into()
.unwrap();
let reshape207_out1 = linear157_out1.reshape(concat183_out1);
let linear158_out1 = self.linear158.forward(layernormalization87_out1.clone());
let reshape208_out1 = linear158_out1.reshape(concat184_out1);
let linear159_out1 = self.linear159.forward(layernormalization87_out1);
let reshape209_out1 = linear159_out1.reshape(concat185_out1);
let transpose174_out1 = reshape209_out1.permute([0, 2, 1, 3]);
let linear160_out1 = self.linear160.forward(slice1_out1);
let reshape210_out1 = linear160_out1.reshape([1, -1, 8, 128]);
let constant432_out1 = self.constant432.val();
let add131_out1 = reshape207_out1
.clone()
.add((constant432_out1).unsqueeze_dims(&[0isize, 1isize]));
let transpose175_out1 = add131_out1.permute([0, 2, 1, 3]);
let constant433_out1 = self.constant433.val();
let add132_out1 = reshape207_out1
.add((constant433_out1).unsqueeze_dims(&[0isize, 1isize]));
let transpose176_out1 = add132_out1.permute([0, 2, 1, 3]);
let transpose177_out1 = reshape210_out1.permute([0, 2, 3, 1]);
let matmul212_out1 = transpose176_out1.matmul(transpose177_out1);
let shape182_out1: [i64; 4] = {
let axes = &matmul212_out1.clone().dims()[0..4];
let mut output = [0i64; 4];
for i in 0..4 {
output[i] = axes[i] as i64;
}
output
};
let gather146_out1 = shape182_out1[0] as i64;
let gather147_out1 = shape182_out1[1] as i64;
let gather148_out1 = shape182_out1[2] as i64;
let gather149_out1 = shape182_out1[3] as i64;
let pad35_out1 = matmul212_out1
.pad(
[(0usize, 0usize), (0usize, 0usize), (0usize, 0usize), (1usize, 0usize)],
burn::tensor::ops::PadMode::Constant(0f32),
);
let unsqueeze316_out1 = [gather146_out1 as i64];
let unsqueeze317_out1 = [gather147_out1 as i64];
let unsqueeze318_out1 = [gather148_out1 as i64];
let constant2322_out1: [i64; 1] = [-1i64];
let concat187_out1: [i64; 4usize] = [
&unsqueeze316_out1[..],
&unsqueeze317_out1[..],
&constant2322_out1[..],
&unsqueeze318_out1[..],
]
.concat()
.try_into()
.unwrap();
let reshape213_out1 = pad35_out1.reshape(concat187_out1);
let slice90_out1 = reshape213_out1.slice(s![.., .., 1.., ..]);
let unsqueeze319_out1 = [gather146_out1 as i64];
let unsqueeze320_out1 = [gather147_out1 as i64];
let unsqueeze321_out1 = [gather148_out1 as i64];
let unsqueeze322_out1 = [gather149_out1 as i64];
let concat188_out1: [i64; 4usize] = [
&unsqueeze319_out1[..],
&unsqueeze320_out1[..],
&unsqueeze321_out1[..],
&unsqueeze322_out1[..],
]
.concat()
.try_into()
.unwrap();
let reshape214_out1 = slice90_out1.reshape(concat188_out1);
let transpose179_out1 = reshape208_out1.permute([0, 2, 3, 1]);
let matmul213_out1 = transpose175_out1.matmul(transpose179_out1);
let shape186_out1: [i64; 4] = {
let axes = &matmul213_out1.clone().dims()[0..4];
let mut output = [0i64; 4];
for i in 0..4 {
output[i] = axes[i] as i64;
}
output
};
let slice91_out1: [i64; 1] = shape186_out1[3..4].try_into().unwrap();
let squeeze18_out1 = slice91_out1[0] as i64;
let unsqueeze323_out1 = [squeeze18_out1 as i64];
let slice92_out1 = reshape214_out1
.slice(s![.., .., .., 0..unsqueeze323_out1[0]]);
let add133_out1 = matmul213_out1.add(slice92_out1);
let constant2340_out1 = self.constant2340.val();
let div21_out1 = add133_out1
.div((constant2340_out1).unsqueeze_dims(&[0isize, 1isize, 2isize]));
let shape187_out1: [i64; 4] = {
let axes = &transpose174_out1.clone().dims()[0..4];
let mut output = [0i64; 4];
for i in 0..4 {
output[i] = axes[i] as i64;
}
output
};
let gather150_out1 = shape187_out1[0] as i64;
let constant2342_out1 = -10000f32;
let where70_out1 = div21_out1
.mask_fill(unsqueeze118_out1.clone(), constant2342_out1);
let softmax18_out1 = burn::tensor::activation::softmax(where70_out1, 3);
let constant2343_out1 = 0f32;
let where71_out1 = softmax18_out1
.mask_fill(unsqueeze118_out1, constant2343_out1);
let matmul214_out1 = where71_out1.matmul(transpose174_out1);
let transpose180_out1 = matmul214_out1.permute([0, 2, 1, 3]);
let unsqueeze324_out1 = [gather150_out1 as i64];
let constant2346_out1: [i64; 1] = [1024i64];
let constant2345_out1: [i64; 1] = [-1i64];
let concat189_out1: [i64; 3usize] = [
&unsqueeze324_out1[..],
&constant2345_out1[..],
&constant2346_out1[..],
]
.concat()
.try_into()
.unwrap();
let reshape215_out1 = transpose180_out1.reshape(concat189_out1);
let linear161_out1 = self.linear161.forward(reshape215_out1);
let add134_out1 = add130_out1.add(linear161_out1);
add134_out1
}
}
#[derive(Module, Debug)]
pub struct Submodule23 {
layernormalization88: LayerNorm,
conv1d52: Conv1d,
conv1d53: Conv1d,
conv1d54: Conv1d,
layernormalization89: LayerNorm,
linear162: Linear,
linear163: Linear,
constant2356: burn::module::Param<Tensor<1>>,
layernormalization90: LayerNorm,
layernormalization91: LayerNorm,
linear164: Linear,
linear165: Linear,
constant2357: burn::module::Param<Tensor<1>>,
layernormalization92: LayerNorm,
linear166: Linear,
linear167: Linear,
linear168: Linear,
linear169: Linear,
constant456: burn::module::Param<Tensor<2>>,
constant457: burn::module::Param<Tensor<2>>,
constant2403: burn::module::Param<Tensor<1>>,
linear170: Linear,
#[module(skip)]
device: Device,
}
impl Submodule23 {
#[allow(unused_variables)]
pub fn new(device: &Device) -> Self {
let layernormalization88 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let conv1d52 = Conv1dConfig::new(1024, 2048, 1)
.with_stride(1)
.with_padding(PaddingConfig1d::Valid)
.with_dilation(1)
.with_groups(1)
.with_bias(true)
.init(device);
let conv1d53 = Conv1dConfig::new(1024, 1024, 9)
.with_stride(1)
.with_padding(PaddingConfig1d::Valid)
.with_dilation(1)
.with_groups(1024)
.with_bias(true)
.init(device);
let conv1d54 = Conv1dConfig::new(1024, 1024, 1)
.with_stride(1)
.with_padding(PaddingConfig1d::Valid)
.with_dilation(1)
.with_groups(1)
.with_bias(true)
.init(device);
let layernormalization89 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let linear162 = LinearConfig::new(1024, 4096).with_bias(true).init(device);
let linear163 = LinearConfig::new(4096, 1024).with_bias(true).init(device);
let constant2356: 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([0.5f64]),
(device, burn::tensor::DType::F32),
),
device.clone(),
false,
[1].into(),
);
let layernormalization90 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let layernormalization91 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let linear164 = LinearConfig::new(1024, 4096).with_bias(true).init(device);
let linear165 = LinearConfig::new(4096, 1024).with_bias(true).init(device);
let constant2357: 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([0.5f64]),
(device, burn::tensor::DType::F32),
),
device.clone(),
false,
[1].into(),
);
let layernormalization92 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let linear166 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
let linear167 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
let linear168 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
let linear169 = LinearConfig::new(1024, 1024).with_bias(false).init(device);
let constant456: burn::module::Param<Tensor<2>> = burn::module::Param::uninitialized(
burn::module::ParamId::new(),
move |device, _require_grad| Tensor::<
2,
>::zeros([8, 128], (device, burn::tensor::DType::F32)),
device.clone(),
false,
[8, 128].into(),
);
let constant457: burn::module::Param<Tensor<2>> = burn::module::Param::uninitialized(
burn::module::ParamId::new(),
move |device, _require_grad| Tensor::<
2,
>::zeros([8, 128], (device, burn::tensor::DType::F32)),
device.clone(),
false,
[8, 128].into(),
);
let constant2403: 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([11.313708305358887f64]),
(device, burn::tensor::DType::F32),
),
device.clone(),
false,
[1].into(),
);
let linear170 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
Self {
layernormalization88,
conv1d52,
conv1d53,
conv1d54,
layernormalization89,
linear162,
linear163,
constant2356,
layernormalization90,
layernormalization91,
linear164,
linear165,
constant2357,
layernormalization92,
linear166,
linear167,
linear168,
linear169,
constant456,
constant457,
constant2403,
linear170,
device: device.clone(),
}
}
#[allow(clippy::let_and_return, clippy::approx_constant)]
pub fn forward(
&self,
add134_out1: Tensor<3>,
unsqueeze120_out1: Tensor<3, Bool>,
slice1_out1: Tensor<3>,
unsqueeze118_out1: Tensor<4, Bool>,
) -> Tensor<3> {
let layernormalization88_out1 = {
let dtype = add134_out1.clone().dtype();
self.layernormalization88
.forward(add134_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let transpose181_out1 = layernormalization88_out1.permute([0, 2, 1]);
let conv1d52_out1 = self.conv1d52.forward(transpose181_out1);
let split_tensors = conv1d52_out1.split_with_sizes([1024, 1024].into(), 1);
let [split18_out1, split18_out2] = split_tensors.try_into().unwrap();
let sigmoid70_out1 = burn::tensor::activation::sigmoid(split18_out2);
let mul151_out1 = split18_out1.mul(sigmoid70_out1);
let constant2347_out1 = 0f32;
let where72_out1 = mul151_out1.mask_fill(unsqueeze120_out1, constant2347_out1);
let pad36_out1 = where72_out1
.pad(
[(0usize, 0usize), (0usize, 0usize), (4usize, 4usize)],
burn::tensor::ops::PadMode::Constant(0f32),
);
let conv1d53_out1 = self.conv1d53.forward(pad36_out1);
let sigmoid71_out1 = burn::tensor::activation::sigmoid(conv1d53_out1.clone());
let mul152_out1 = conv1d53_out1.mul(sigmoid71_out1);
let conv1d54_out1 = self.conv1d54.forward(mul152_out1);
let transpose183_out1 = conv1d54_out1.permute([0, 2, 1]);
let add135_out1 = add134_out1.add(transpose183_out1);
let layernormalization89_out1 = {
let dtype = add135_out1.clone().dtype();
self.layernormalization89
.forward(add135_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let linear162_out1 = self.linear162.forward(layernormalization89_out1);
let sigmoid72_out1 = burn::tensor::activation::sigmoid(linear162_out1.clone());
let mul153_out1 = linear162_out1.mul(sigmoid72_out1);
let linear163_out1 = self.linear163.forward(mul153_out1);
let constant2356_out1 = self.constant2356.val();
let mul154_out1 = linear163_out1
.mul((constant2356_out1).unsqueeze_dims(&[0isize, 1isize]));
let add136_out1 = add135_out1.add(mul154_out1);
let layernormalization90_out1 = {
let dtype = add136_out1.dtype();
self.layernormalization90
.forward(add136_out1.cast(burn::tensor::DType::F32))
.cast(dtype)
};
let layernormalization91_out1 = {
let dtype = layernormalization90_out1.clone().dtype();
self.layernormalization91
.forward(
layernormalization90_out1.clone().cast(burn::tensor::DType::F32),
)
.cast(dtype)
};
let linear164_out1 = self.linear164.forward(layernormalization91_out1);
let sigmoid73_out1 = burn::tensor::activation::sigmoid(linear164_out1.clone());
let mul155_out1 = linear164_out1.mul(sigmoid73_out1);
let linear165_out1 = self.linear165.forward(mul155_out1);
let constant2357_out1 = self.constant2357.val();
let mul156_out1 = linear165_out1
.mul((constant2357_out1).unsqueeze_dims(&[0isize, 1isize]));
let add137_out1 = layernormalization90_out1.add(mul156_out1);
let layernormalization92_out1 = {
let dtype = add137_out1.clone().dtype();
self.layernormalization92
.forward(add137_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let shape188_out1: [i64; 3] = {
let axes = &layernormalization92_out1.clone().dims()[0..3];
let mut output = [0i64; 3];
for i in 0..3 {
output[i] = axes[i] as i64;
}
output
};
let gather151_out1 = shape188_out1[0] as i64;
let linear166_out1 = self.linear166.forward(layernormalization92_out1.clone());
let unsqueeze325_out1 = [gather151_out1 as i64];
let constant2362_out1: [i64; 1] = [128i64];
let constant2360_out1: [i64; 1] = [-1i64];
let constant2361_out1: [i64; 1] = [8i64];
let concat191_out1: [i64; 4usize] = [
&unsqueeze325_out1[..],
&constant2360_out1[..],
&constant2361_out1[..],
&constant2362_out1[..],
]
.concat()
.try_into()
.unwrap();
let unsqueeze326_out1 = [gather151_out1 as i64];
let constant2366_out1: [i64; 1] = [128i64];
let constant2364_out1: [i64; 1] = [-1i64];
let constant2365_out1: [i64; 1] = [8i64];
let concat192_out1: [i64; 4usize] = [
&unsqueeze326_out1[..],
&constant2364_out1[..],
&constant2365_out1[..],
&constant2366_out1[..],
]
.concat()
.try_into()
.unwrap();
let unsqueeze327_out1 = [gather151_out1 as i64];
let constant2370_out1: [i64; 1] = [128i64];
let constant2368_out1: [i64; 1] = [-1i64];
let constant2369_out1: [i64; 1] = [8i64];
let concat193_out1: [i64; 4usize] = [
&unsqueeze327_out1[..],
&constant2368_out1[..],
&constant2369_out1[..],
&constant2370_out1[..],
]
.concat()
.try_into()
.unwrap();
let reshape218_out1 = linear166_out1.reshape(concat191_out1);
let linear167_out1 = self.linear167.forward(layernormalization92_out1.clone());
let reshape219_out1 = linear167_out1.reshape(concat192_out1);
let linear168_out1 = self.linear168.forward(layernormalization92_out1);
let reshape220_out1 = linear168_out1.reshape(concat193_out1);
let transpose184_out1 = reshape220_out1.permute([0, 2, 1, 3]);
let linear169_out1 = self.linear169.forward(slice1_out1);
let reshape221_out1 = linear169_out1.reshape([1, -1, 8, 128]);
let constant456_out1 = self.constant456.val();
let add138_out1 = reshape218_out1
.clone()
.add((constant456_out1).unsqueeze_dims(&[0isize, 1isize]));
let transpose185_out1 = add138_out1.permute([0, 2, 1, 3]);
let constant457_out1 = self.constant457.val();
let add139_out1 = reshape218_out1
.add((constant457_out1).unsqueeze_dims(&[0isize, 1isize]));
let transpose186_out1 = add139_out1.permute([0, 2, 1, 3]);
let transpose187_out1 = reshape221_out1.permute([0, 2, 3, 1]);
let matmul224_out1 = transpose186_out1.matmul(transpose187_out1);
let shape189_out1: [i64; 4] = {
let axes = &matmul224_out1.clone().dims()[0..4];
let mut output = [0i64; 4];
for i in 0..4 {
output[i] = axes[i] as i64;
}
output
};
let gather152_out1 = shape189_out1[0] as i64;
let gather153_out1 = shape189_out1[1] as i64;
let gather154_out1 = shape189_out1[2] as i64;
let gather155_out1 = shape189_out1[3] as i64;
let pad37_out1 = matmul224_out1
.pad(
[(0usize, 0usize), (0usize, 0usize), (0usize, 0usize), (1usize, 0usize)],
burn::tensor::ops::PadMode::Constant(0f32),
);
let unsqueeze328_out1 = [gather152_out1 as i64];
let unsqueeze329_out1 = [gather153_out1 as i64];
let unsqueeze330_out1 = [gather154_out1 as i64];
let constant2385_out1: [i64; 1] = [-1i64];
let concat195_out1: [i64; 4usize] = [
&unsqueeze328_out1[..],
&unsqueeze329_out1[..],
&constant2385_out1[..],
&unsqueeze330_out1[..],
]
.concat()
.try_into()
.unwrap();
let reshape224_out1 = pad37_out1.reshape(concat195_out1);
let slice95_out1 = reshape224_out1.slice(s![.., .., 1.., ..]);
let unsqueeze331_out1 = [gather152_out1 as i64];
let unsqueeze332_out1 = [gather153_out1 as i64];
let unsqueeze333_out1 = [gather154_out1 as i64];
let unsqueeze334_out1 = [gather155_out1 as i64];
let concat196_out1: [i64; 4usize] = [
&unsqueeze331_out1[..],
&unsqueeze332_out1[..],
&unsqueeze333_out1[..],
&unsqueeze334_out1[..],
]
.concat()
.try_into()
.unwrap();
let reshape225_out1 = slice95_out1.reshape(concat196_out1);
let transpose189_out1 = reshape219_out1.permute([0, 2, 3, 1]);
let matmul225_out1 = transpose185_out1.matmul(transpose189_out1);
let shape193_out1: [i64; 4] = {
let axes = &matmul225_out1.clone().dims()[0..4];
let mut output = [0i64; 4];
for i in 0..4 {
output[i] = axes[i] as i64;
}
output
};
let slice96_out1: [i64; 1] = shape193_out1[3..4].try_into().unwrap();
let squeeze19_out1 = slice96_out1[0] as i64;
let unsqueeze335_out1 = [squeeze19_out1 as i64];
let slice97_out1 = reshape225_out1
.slice(s![.., .., .., 0..unsqueeze335_out1[0]]);
let add140_out1 = matmul225_out1.add(slice97_out1);
let constant2403_out1 = self.constant2403.val();
let div22_out1 = add140_out1
.div((constant2403_out1).unsqueeze_dims(&[0isize, 1isize, 2isize]));
let shape194_out1: [i64; 4] = {
let axes = &transpose184_out1.clone().dims()[0..4];
let mut output = [0i64; 4];
for i in 0..4 {
output[i] = axes[i] as i64;
}
output
};
let gather156_out1 = shape194_out1[0] as i64;
let constant2405_out1 = -10000f32;
let where73_out1 = div22_out1
.mask_fill(unsqueeze118_out1.clone(), constant2405_out1);
let softmax19_out1 = burn::tensor::activation::softmax(where73_out1, 3);
let constant2406_out1 = 0f32;
let where74_out1 = softmax19_out1
.mask_fill(unsqueeze118_out1, constant2406_out1);
let matmul226_out1 = where74_out1.matmul(transpose184_out1);
let transpose190_out1 = matmul226_out1.permute([0, 2, 1, 3]);
let unsqueeze336_out1 = [gather156_out1 as i64];
let constant2409_out1: [i64; 1] = [1024i64];
let constant2408_out1: [i64; 1] = [-1i64];
let concat197_out1: [i64; 3usize] = [
&unsqueeze336_out1[..],
&constant2408_out1[..],
&constant2409_out1[..],
]
.concat()
.try_into()
.unwrap();
let reshape226_out1 = transpose190_out1.reshape(concat197_out1);
let linear170_out1 = self.linear170.forward(reshape226_out1);
let add141_out1 = add137_out1.add(linear170_out1);
add141_out1
}
}
#[derive(Module, Debug)]
pub struct Submodule24 {
layernormalization93: LayerNorm,
conv1d55: Conv1d,
conv1d56: Conv1d,
conv1d57: Conv1d,
layernormalization94: LayerNorm,
linear171: Linear,
linear172: Linear,
constant2419: burn::module::Param<Tensor<1>>,
layernormalization95: LayerNorm,
layernormalization96: LayerNorm,
linear173: Linear,
linear174: Linear,
constant2420: burn::module::Param<Tensor<1>>,
layernormalization97: LayerNorm,
linear175: Linear,
linear176: Linear,
linear177: Linear,
linear178: Linear,
constant480: burn::module::Param<Tensor<2>>,
constant481: burn::module::Param<Tensor<2>>,
constant2466: burn::module::Param<Tensor<1>>,
linear179: Linear,
#[module(skip)]
device: Device,
}
impl Submodule24 {
#[allow(unused_variables)]
pub fn new(device: &Device) -> Self {
let layernormalization93 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let conv1d55 = Conv1dConfig::new(1024, 2048, 1)
.with_stride(1)
.with_padding(PaddingConfig1d::Valid)
.with_dilation(1)
.with_groups(1)
.with_bias(true)
.init(device);
let conv1d56 = Conv1dConfig::new(1024, 1024, 9)
.with_stride(1)
.with_padding(PaddingConfig1d::Valid)
.with_dilation(1)
.with_groups(1024)
.with_bias(true)
.init(device);
let conv1d57 = Conv1dConfig::new(1024, 1024, 1)
.with_stride(1)
.with_padding(PaddingConfig1d::Valid)
.with_dilation(1)
.with_groups(1)
.with_bias(true)
.init(device);
let layernormalization94 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let linear171 = LinearConfig::new(1024, 4096).with_bias(true).init(device);
let linear172 = LinearConfig::new(4096, 1024).with_bias(true).init(device);
let constant2419: 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([0.5f64]),
(device, burn::tensor::DType::F32),
),
device.clone(),
false,
[1].into(),
);
let layernormalization95 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let layernormalization96 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let linear173 = LinearConfig::new(1024, 4096).with_bias(true).init(device);
let linear174 = LinearConfig::new(4096, 1024).with_bias(true).init(device);
let constant2420: 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([0.5f64]),
(device, burn::tensor::DType::F32),
),
device.clone(),
false,
[1].into(),
);
let layernormalization97 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let linear175 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
let linear176 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
let linear177 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
let linear178 = LinearConfig::new(1024, 1024).with_bias(false).init(device);
let constant480: burn::module::Param<Tensor<2>> = burn::module::Param::uninitialized(
burn::module::ParamId::new(),
move |device, _require_grad| Tensor::<
2,
>::zeros([8, 128], (device, burn::tensor::DType::F32)),
device.clone(),
false,
[8, 128].into(),
);
let constant481: burn::module::Param<Tensor<2>> = burn::module::Param::uninitialized(
burn::module::ParamId::new(),
move |device, _require_grad| Tensor::<
2,
>::zeros([8, 128], (device, burn::tensor::DType::F32)),
device.clone(),
false,
[8, 128].into(),
);
let constant2466: 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([11.313708305358887f64]),
(device, burn::tensor::DType::F32),
),
device.clone(),
false,
[1].into(),
);
let linear179 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
Self {
layernormalization93,
conv1d55,
conv1d56,
conv1d57,
layernormalization94,
linear171,
linear172,
constant2419,
layernormalization95,
layernormalization96,
linear173,
linear174,
constant2420,
layernormalization97,
linear175,
linear176,
linear177,
linear178,
constant480,
constant481,
constant2466,
linear179,
device: device.clone(),
}
}
#[allow(clippy::let_and_return, clippy::approx_constant)]
pub fn forward(
&self,
add141_out1: Tensor<3>,
unsqueeze120_out1: Tensor<3, Bool>,
slice1_out1: Tensor<3>,
unsqueeze118_out1: Tensor<4, Bool>,
) -> Tensor<3> {
let layernormalization93_out1 = {
let dtype = add141_out1.clone().dtype();
self.layernormalization93
.forward(add141_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let transpose191_out1 = layernormalization93_out1.permute([0, 2, 1]);
let conv1d55_out1 = self.conv1d55.forward(transpose191_out1);
let split_tensors = conv1d55_out1.split_with_sizes([1024, 1024].into(), 1);
let [split19_out1, split19_out2] = split_tensors.try_into().unwrap();
let sigmoid74_out1 = burn::tensor::activation::sigmoid(split19_out2);
let mul157_out1 = split19_out1.mul(sigmoid74_out1);
let constant2410_out1 = 0f32;
let where75_out1 = mul157_out1.mask_fill(unsqueeze120_out1, constant2410_out1);
let pad38_out1 = where75_out1
.pad(
[(0usize, 0usize), (0usize, 0usize), (4usize, 4usize)],
burn::tensor::ops::PadMode::Constant(0f32),
);
let conv1d56_out1 = self.conv1d56.forward(pad38_out1);
let sigmoid75_out1 = burn::tensor::activation::sigmoid(conv1d56_out1.clone());
let mul158_out1 = conv1d56_out1.mul(sigmoid75_out1);
let conv1d57_out1 = self.conv1d57.forward(mul158_out1);
let transpose193_out1 = conv1d57_out1.permute([0, 2, 1]);
let add142_out1 = add141_out1.add(transpose193_out1);
let layernormalization94_out1 = {
let dtype = add142_out1.clone().dtype();
self.layernormalization94
.forward(add142_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let linear171_out1 = self.linear171.forward(layernormalization94_out1);
let sigmoid76_out1 = burn::tensor::activation::sigmoid(linear171_out1.clone());
let mul159_out1 = linear171_out1.mul(sigmoid76_out1);
let linear172_out1 = self.linear172.forward(mul159_out1);
let constant2419_out1 = self.constant2419.val();
let mul160_out1 = linear172_out1
.mul((constant2419_out1).unsqueeze_dims(&[0isize, 1isize]));
let add143_out1 = add142_out1.add(mul160_out1);
let layernormalization95_out1 = {
let dtype = add143_out1.dtype();
self.layernormalization95
.forward(add143_out1.cast(burn::tensor::DType::F32))
.cast(dtype)
};
let layernormalization96_out1 = {
let dtype = layernormalization95_out1.clone().dtype();
self.layernormalization96
.forward(
layernormalization95_out1.clone().cast(burn::tensor::DType::F32),
)
.cast(dtype)
};
let linear173_out1 = self.linear173.forward(layernormalization96_out1);
let sigmoid77_out1 = burn::tensor::activation::sigmoid(linear173_out1.clone());
let mul161_out1 = linear173_out1.mul(sigmoid77_out1);
let linear174_out1 = self.linear174.forward(mul161_out1);
let constant2420_out1 = self.constant2420.val();
let mul162_out1 = linear174_out1
.mul((constant2420_out1).unsqueeze_dims(&[0isize, 1isize]));
let add144_out1 = layernormalization95_out1.add(mul162_out1);
let layernormalization97_out1 = {
let dtype = add144_out1.clone().dtype();
self.layernormalization97
.forward(add144_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let shape195_out1: [i64; 3] = {
let axes = &layernormalization97_out1.clone().dims()[0..3];
let mut output = [0i64; 3];
for i in 0..3 {
output[i] = axes[i] as i64;
}
output
};
let gather157_out1 = shape195_out1[0] as i64;
let linear175_out1 = self.linear175.forward(layernormalization97_out1.clone());
let unsqueeze337_out1 = [gather157_out1 as i64];
let constant2425_out1: [i64; 1] = [128i64];
let constant2423_out1: [i64; 1] = [-1i64];
let constant2424_out1: [i64; 1] = [8i64];
let concat199_out1: [i64; 4usize] = [
&unsqueeze337_out1[..],
&constant2423_out1[..],
&constant2424_out1[..],
&constant2425_out1[..],
]
.concat()
.try_into()
.unwrap();
let unsqueeze338_out1 = [gather157_out1 as i64];
let constant2429_out1: [i64; 1] = [128i64];
let constant2427_out1: [i64; 1] = [-1i64];
let constant2428_out1: [i64; 1] = [8i64];
let concat200_out1: [i64; 4usize] = [
&unsqueeze338_out1[..],
&constant2427_out1[..],
&constant2428_out1[..],
&constant2429_out1[..],
]
.concat()
.try_into()
.unwrap();
let unsqueeze339_out1 = [gather157_out1 as i64];
let constant2433_out1: [i64; 1] = [128i64];
let constant2431_out1: [i64; 1] = [-1i64];
let constant2432_out1: [i64; 1] = [8i64];
let concat201_out1: [i64; 4usize] = [
&unsqueeze339_out1[..],
&constant2431_out1[..],
&constant2432_out1[..],
&constant2433_out1[..],
]
.concat()
.try_into()
.unwrap();
let reshape229_out1 = linear175_out1.reshape(concat199_out1);
let linear176_out1 = self.linear176.forward(layernormalization97_out1.clone());
let reshape230_out1 = linear176_out1.reshape(concat200_out1);
let linear177_out1 = self.linear177.forward(layernormalization97_out1);
let reshape231_out1 = linear177_out1.reshape(concat201_out1);
let transpose194_out1 = reshape231_out1.permute([0, 2, 1, 3]);
let linear178_out1 = self.linear178.forward(slice1_out1);
let reshape232_out1 = linear178_out1.reshape([1, -1, 8, 128]);
let constant480_out1 = self.constant480.val();
let add145_out1 = reshape229_out1
.clone()
.add((constant480_out1).unsqueeze_dims(&[0isize, 1isize]));
let transpose195_out1 = add145_out1.permute([0, 2, 1, 3]);
let constant481_out1 = self.constant481.val();
let add146_out1 = reshape229_out1
.add((constant481_out1).unsqueeze_dims(&[0isize, 1isize]));
let transpose196_out1 = add146_out1.permute([0, 2, 1, 3]);
let transpose197_out1 = reshape232_out1.permute([0, 2, 3, 1]);
let matmul236_out1 = transpose196_out1.matmul(transpose197_out1);
let shape196_out1: [i64; 4] = {
let axes = &matmul236_out1.clone().dims()[0..4];
let mut output = [0i64; 4];
for i in 0..4 {
output[i] = axes[i] as i64;
}
output
};
let gather158_out1 = shape196_out1[0] as i64;
let gather159_out1 = shape196_out1[1] as i64;
let gather160_out1 = shape196_out1[2] as i64;
let gather161_out1 = shape196_out1[3] as i64;
let pad39_out1 = matmul236_out1
.pad(
[(0usize, 0usize), (0usize, 0usize), (0usize, 0usize), (1usize, 0usize)],
burn::tensor::ops::PadMode::Constant(0f32),
);
let unsqueeze340_out1 = [gather158_out1 as i64];
let unsqueeze341_out1 = [gather159_out1 as i64];
let unsqueeze342_out1 = [gather160_out1 as i64];
let constant2448_out1: [i64; 1] = [-1i64];
let concat203_out1: [i64; 4usize] = [
&unsqueeze340_out1[..],
&unsqueeze341_out1[..],
&constant2448_out1[..],
&unsqueeze342_out1[..],
]
.concat()
.try_into()
.unwrap();
let reshape235_out1 = pad39_out1.reshape(concat203_out1);
let slice100_out1 = reshape235_out1.slice(s![.., .., 1.., ..]);
let unsqueeze343_out1 = [gather158_out1 as i64];
let unsqueeze344_out1 = [gather159_out1 as i64];
let unsqueeze345_out1 = [gather160_out1 as i64];
let unsqueeze346_out1 = [gather161_out1 as i64];
let concat204_out1: [i64; 4usize] = [
&unsqueeze343_out1[..],
&unsqueeze344_out1[..],
&unsqueeze345_out1[..],
&unsqueeze346_out1[..],
]
.concat()
.try_into()
.unwrap();
let reshape236_out1 = slice100_out1.reshape(concat204_out1);
let transpose199_out1 = reshape230_out1.permute([0, 2, 3, 1]);
let matmul237_out1 = transpose195_out1.matmul(transpose199_out1);
let shape200_out1: [i64; 4] = {
let axes = &matmul237_out1.clone().dims()[0..4];
let mut output = [0i64; 4];
for i in 0..4 {
output[i] = axes[i] as i64;
}
output
};
let slice101_out1: [i64; 1] = shape200_out1[3..4].try_into().unwrap();
let squeeze20_out1 = slice101_out1[0] as i64;
let unsqueeze347_out1 = [squeeze20_out1 as i64];
let slice102_out1 = reshape236_out1
.slice(s![.., .., .., 0..unsqueeze347_out1[0]]);
let add147_out1 = matmul237_out1.add(slice102_out1);
let constant2466_out1 = self.constant2466.val();
let div23_out1 = add147_out1
.div((constant2466_out1).unsqueeze_dims(&[0isize, 1isize, 2isize]));
let shape201_out1: [i64; 4] = {
let axes = &transpose194_out1.clone().dims()[0..4];
let mut output = [0i64; 4];
for i in 0..4 {
output[i] = axes[i] as i64;
}
output
};
let gather162_out1 = shape201_out1[0] as i64;
let constant2468_out1 = -10000f32;
let where76_out1 = div23_out1
.mask_fill(unsqueeze118_out1.clone(), constant2468_out1);
let softmax20_out1 = burn::tensor::activation::softmax(where76_out1, 3);
let constant2469_out1 = 0f32;
let where77_out1 = softmax20_out1
.mask_fill(unsqueeze118_out1, constant2469_out1);
let matmul238_out1 = where77_out1.matmul(transpose194_out1);
let transpose200_out1 = matmul238_out1.permute([0, 2, 1, 3]);
let unsqueeze348_out1 = [gather162_out1 as i64];
let constant2472_out1: [i64; 1] = [1024i64];
let constant2471_out1: [i64; 1] = [-1i64];
let concat205_out1: [i64; 3usize] = [
&unsqueeze348_out1[..],
&constant2471_out1[..],
&constant2472_out1[..],
]
.concat()
.try_into()
.unwrap();
let reshape237_out1 = transpose200_out1.reshape(concat205_out1);
let linear179_out1 = self.linear179.forward(reshape237_out1);
let add148_out1 = add144_out1.add(linear179_out1);
add148_out1
}
}
#[derive(Module, Debug)]
pub struct Submodule25 {
layernormalization98: LayerNorm,
conv1d58: Conv1d,
conv1d59: Conv1d,
conv1d60: Conv1d,
layernormalization99: LayerNorm,
linear180: Linear,
linear181: Linear,
constant2482: burn::module::Param<Tensor<1>>,
layernormalization100: LayerNorm,
layernormalization101: LayerNorm,
linear182: Linear,
linear183: Linear,
constant2483: burn::module::Param<Tensor<1>>,
layernormalization102: LayerNorm,
linear184: Linear,
linear185: Linear,
linear186: Linear,
linear187: Linear,
constant504: burn::module::Param<Tensor<2>>,
constant505: burn::module::Param<Tensor<2>>,
constant2529: burn::module::Param<Tensor<1>>,
linear188: Linear,
#[module(skip)]
device: Device,
}
impl Submodule25 {
#[allow(unused_variables)]
pub fn new(device: &Device) -> Self {
let layernormalization98 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let conv1d58 = Conv1dConfig::new(1024, 2048, 1)
.with_stride(1)
.with_padding(PaddingConfig1d::Valid)
.with_dilation(1)
.with_groups(1)
.with_bias(true)
.init(device);
let conv1d59 = Conv1dConfig::new(1024, 1024, 9)
.with_stride(1)
.with_padding(PaddingConfig1d::Valid)
.with_dilation(1)
.with_groups(1024)
.with_bias(true)
.init(device);
let conv1d60 = Conv1dConfig::new(1024, 1024, 1)
.with_stride(1)
.with_padding(PaddingConfig1d::Valid)
.with_dilation(1)
.with_groups(1)
.with_bias(true)
.init(device);
let layernormalization99 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let linear180 = LinearConfig::new(1024, 4096).with_bias(true).init(device);
let linear181 = LinearConfig::new(4096, 1024).with_bias(true).init(device);
let constant2482: 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([0.5f64]),
(device, burn::tensor::DType::F32),
),
device.clone(),
false,
[1].into(),
);
let layernormalization100 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let layernormalization101 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let linear182 = LinearConfig::new(1024, 4096).with_bias(true).init(device);
let linear183 = LinearConfig::new(4096, 1024).with_bias(true).init(device);
let constant2483: 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([0.5f64]),
(device, burn::tensor::DType::F32),
),
device.clone(),
false,
[1].into(),
);
let layernormalization102 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let linear184 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
let linear185 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
let linear186 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
let linear187 = LinearConfig::new(1024, 1024).with_bias(false).init(device);
let constant504: burn::module::Param<Tensor<2>> = burn::module::Param::uninitialized(
burn::module::ParamId::new(),
move |device, _require_grad| Tensor::<
2,
>::zeros([8, 128], (device, burn::tensor::DType::F32)),
device.clone(),
false,
[8, 128].into(),
);
let constant505: burn::module::Param<Tensor<2>> = burn::module::Param::uninitialized(
burn::module::ParamId::new(),
move |device, _require_grad| Tensor::<
2,
>::zeros([8, 128], (device, burn::tensor::DType::F32)),
device.clone(),
false,
[8, 128].into(),
);
let constant2529: 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([11.313708305358887f64]),
(device, burn::tensor::DType::F32),
),
device.clone(),
false,
[1].into(),
);
let linear188 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
Self {
layernormalization98,
conv1d58,
conv1d59,
conv1d60,
layernormalization99,
linear180,
linear181,
constant2482,
layernormalization100,
layernormalization101,
linear182,
linear183,
constant2483,
layernormalization102,
linear184,
linear185,
linear186,
linear187,
constant504,
constant505,
constant2529,
linear188,
device: device.clone(),
}
}
#[allow(clippy::let_and_return, clippy::approx_constant)]
pub fn forward(
&self,
add148_out1: Tensor<3>,
unsqueeze120_out1: Tensor<3, Bool>,
slice1_out1: Tensor<3>,
unsqueeze118_out1: Tensor<4, Bool>,
) -> Tensor<3> {
let layernormalization98_out1 = {
let dtype = add148_out1.clone().dtype();
self.layernormalization98
.forward(add148_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let transpose201_out1 = layernormalization98_out1.permute([0, 2, 1]);
let conv1d58_out1 = self.conv1d58.forward(transpose201_out1);
let split_tensors = conv1d58_out1.split_with_sizes([1024, 1024].into(), 1);
let [split20_out1, split20_out2] = split_tensors.try_into().unwrap();
let sigmoid78_out1 = burn::tensor::activation::sigmoid(split20_out2);
let mul163_out1 = split20_out1.mul(sigmoid78_out1);
let constant2473_out1 = 0f32;
let where78_out1 = mul163_out1.mask_fill(unsqueeze120_out1, constant2473_out1);
let pad40_out1 = where78_out1
.pad(
[(0usize, 0usize), (0usize, 0usize), (4usize, 4usize)],
burn::tensor::ops::PadMode::Constant(0f32),
);
let conv1d59_out1 = self.conv1d59.forward(pad40_out1);
let sigmoid79_out1 = burn::tensor::activation::sigmoid(conv1d59_out1.clone());
let mul164_out1 = conv1d59_out1.mul(sigmoid79_out1);
let conv1d60_out1 = self.conv1d60.forward(mul164_out1);
let transpose203_out1 = conv1d60_out1.permute([0, 2, 1]);
let add149_out1 = add148_out1.add(transpose203_out1);
let layernormalization99_out1 = {
let dtype = add149_out1.clone().dtype();
self.layernormalization99
.forward(add149_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let linear180_out1 = self.linear180.forward(layernormalization99_out1);
let sigmoid80_out1 = burn::tensor::activation::sigmoid(linear180_out1.clone());
let mul165_out1 = linear180_out1.mul(sigmoid80_out1);
let linear181_out1 = self.linear181.forward(mul165_out1);
let constant2482_out1 = self.constant2482.val();
let mul166_out1 = linear181_out1
.mul((constant2482_out1).unsqueeze_dims(&[0isize, 1isize]));
let add150_out1 = add149_out1.add(mul166_out1);
let layernormalization100_out1 = {
let dtype = add150_out1.dtype();
self.layernormalization100
.forward(add150_out1.cast(burn::tensor::DType::F32))
.cast(dtype)
};
let layernormalization101_out1 = {
let dtype = layernormalization100_out1.clone().dtype();
self.layernormalization101
.forward(
layernormalization100_out1.clone().cast(burn::tensor::DType::F32),
)
.cast(dtype)
};
let linear182_out1 = self.linear182.forward(layernormalization101_out1);
let sigmoid81_out1 = burn::tensor::activation::sigmoid(linear182_out1.clone());
let mul167_out1 = linear182_out1.mul(sigmoid81_out1);
let linear183_out1 = self.linear183.forward(mul167_out1);
let constant2483_out1 = self.constant2483.val();
let mul168_out1 = linear183_out1
.mul((constant2483_out1).unsqueeze_dims(&[0isize, 1isize]));
let add151_out1 = layernormalization100_out1.add(mul168_out1);
let layernormalization102_out1 = {
let dtype = add151_out1.clone().dtype();
self.layernormalization102
.forward(add151_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let shape202_out1: [i64; 3] = {
let axes = &layernormalization102_out1.clone().dims()[0..3];
let mut output = [0i64; 3];
for i in 0..3 {
output[i] = axes[i] as i64;
}
output
};
let gather163_out1 = shape202_out1[0] as i64;
let linear184_out1 = self.linear184.forward(layernormalization102_out1.clone());
let unsqueeze349_out1 = [gather163_out1 as i64];
let constant2488_out1: [i64; 1] = [128i64];
let constant2486_out1: [i64; 1] = [-1i64];
let constant2487_out1: [i64; 1] = [8i64];
let concat207_out1: [i64; 4usize] = [
&unsqueeze349_out1[..],
&constant2486_out1[..],
&constant2487_out1[..],
&constant2488_out1[..],
]
.concat()
.try_into()
.unwrap();
let unsqueeze350_out1 = [gather163_out1 as i64];
let constant2492_out1: [i64; 1] = [128i64];
let constant2490_out1: [i64; 1] = [-1i64];
let constant2491_out1: [i64; 1] = [8i64];
let concat208_out1: [i64; 4usize] = [
&unsqueeze350_out1[..],
&constant2490_out1[..],
&constant2491_out1[..],
&constant2492_out1[..],
]
.concat()
.try_into()
.unwrap();
let unsqueeze351_out1 = [gather163_out1 as i64];
let constant2496_out1: [i64; 1] = [128i64];
let constant2494_out1: [i64; 1] = [-1i64];
let constant2495_out1: [i64; 1] = [8i64];
let concat209_out1: [i64; 4usize] = [
&unsqueeze351_out1[..],
&constant2494_out1[..],
&constant2495_out1[..],
&constant2496_out1[..],
]
.concat()
.try_into()
.unwrap();
let reshape240_out1 = linear184_out1.reshape(concat207_out1);
let linear185_out1 = self.linear185.forward(layernormalization102_out1.clone());
let reshape241_out1 = linear185_out1.reshape(concat208_out1);
let linear186_out1 = self.linear186.forward(layernormalization102_out1);
let reshape242_out1 = linear186_out1.reshape(concat209_out1);
let transpose204_out1 = reshape242_out1.permute([0, 2, 1, 3]);
let linear187_out1 = self.linear187.forward(slice1_out1);
let reshape243_out1 = linear187_out1.reshape([1, -1, 8, 128]);
let constant504_out1 = self.constant504.val();
let add152_out1 = reshape240_out1
.clone()
.add((constant504_out1).unsqueeze_dims(&[0isize, 1isize]));
let transpose205_out1 = add152_out1.permute([0, 2, 1, 3]);
let constant505_out1 = self.constant505.val();
let add153_out1 = reshape240_out1
.add((constant505_out1).unsqueeze_dims(&[0isize, 1isize]));
let transpose206_out1 = add153_out1.permute([0, 2, 1, 3]);
let transpose207_out1 = reshape243_out1.permute([0, 2, 3, 1]);
let matmul248_out1 = transpose206_out1.matmul(transpose207_out1);
let shape203_out1: [i64; 4] = {
let axes = &matmul248_out1.clone().dims()[0..4];
let mut output = [0i64; 4];
for i in 0..4 {
output[i] = axes[i] as i64;
}
output
};
let gather164_out1 = shape203_out1[0] as i64;
let gather165_out1 = shape203_out1[1] as i64;
let gather166_out1 = shape203_out1[2] as i64;
let gather167_out1 = shape203_out1[3] as i64;
let pad41_out1 = matmul248_out1
.pad(
[(0usize, 0usize), (0usize, 0usize), (0usize, 0usize), (1usize, 0usize)],
burn::tensor::ops::PadMode::Constant(0f32),
);
let unsqueeze352_out1 = [gather164_out1 as i64];
let unsqueeze353_out1 = [gather165_out1 as i64];
let unsqueeze354_out1 = [gather166_out1 as i64];
let constant2511_out1: [i64; 1] = [-1i64];
let concat211_out1: [i64; 4usize] = [
&unsqueeze352_out1[..],
&unsqueeze353_out1[..],
&constant2511_out1[..],
&unsqueeze354_out1[..],
]
.concat()
.try_into()
.unwrap();
let reshape246_out1 = pad41_out1.reshape(concat211_out1);
let slice105_out1 = reshape246_out1.slice(s![.., .., 1.., ..]);
let unsqueeze355_out1 = [gather164_out1 as i64];
let unsqueeze356_out1 = [gather165_out1 as i64];
let unsqueeze357_out1 = [gather166_out1 as i64];
let unsqueeze358_out1 = [gather167_out1 as i64];
let concat212_out1: [i64; 4usize] = [
&unsqueeze355_out1[..],
&unsqueeze356_out1[..],
&unsqueeze357_out1[..],
&unsqueeze358_out1[..],
]
.concat()
.try_into()
.unwrap();
let reshape247_out1 = slice105_out1.reshape(concat212_out1);
let transpose209_out1 = reshape241_out1.permute([0, 2, 3, 1]);
let matmul249_out1 = transpose205_out1.matmul(transpose209_out1);
let shape207_out1: [i64; 4] = {
let axes = &matmul249_out1.clone().dims()[0..4];
let mut output = [0i64; 4];
for i in 0..4 {
output[i] = axes[i] as i64;
}
output
};
let slice106_out1: [i64; 1] = shape207_out1[3..4].try_into().unwrap();
let squeeze21_out1 = slice106_out1[0] as i64;
let unsqueeze359_out1 = [squeeze21_out1 as i64];
let slice107_out1 = reshape247_out1
.slice(s![.., .., .., 0..unsqueeze359_out1[0]]);
let add154_out1 = matmul249_out1.add(slice107_out1);
let constant2529_out1 = self.constant2529.val();
let div24_out1 = add154_out1
.div((constant2529_out1).unsqueeze_dims(&[0isize, 1isize, 2isize]));
let shape208_out1: [i64; 4] = {
let axes = &transpose204_out1.clone().dims()[0..4];
let mut output = [0i64; 4];
for i in 0..4 {
output[i] = axes[i] as i64;
}
output
};
let gather168_out1 = shape208_out1[0] as i64;
let constant2531_out1 = -10000f32;
let where79_out1 = div24_out1
.mask_fill(unsqueeze118_out1.clone(), constant2531_out1);
let softmax21_out1 = burn::tensor::activation::softmax(where79_out1, 3);
let constant2532_out1 = 0f32;
let where80_out1 = softmax21_out1
.mask_fill(unsqueeze118_out1, constant2532_out1);
let matmul250_out1 = where80_out1.matmul(transpose204_out1);
let transpose210_out1 = matmul250_out1.permute([0, 2, 1, 3]);
let unsqueeze360_out1 = [gather168_out1 as i64];
let constant2535_out1: [i64; 1] = [1024i64];
let constant2534_out1: [i64; 1] = [-1i64];
let concat213_out1: [i64; 3usize] = [
&unsqueeze360_out1[..],
&constant2534_out1[..],
&constant2535_out1[..],
]
.concat()
.try_into()
.unwrap();
let reshape248_out1 = transpose210_out1.reshape(concat213_out1);
let linear188_out1 = self.linear188.forward(reshape248_out1);
let add155_out1 = add151_out1.add(linear188_out1);
add155_out1
}
}
#[derive(Module, Debug)]
pub struct Submodule26 {
layernormalization103: LayerNorm,
conv1d61: Conv1d,
conv1d62: Conv1d,
conv1d63: Conv1d,
layernormalization104: LayerNorm,
linear189: Linear,
linear190: Linear,
constant2545: burn::module::Param<Tensor<1>>,
layernormalization105: LayerNorm,
layernormalization106: LayerNorm,
linear191: Linear,
linear192: Linear,
constant2546: burn::module::Param<Tensor<1>>,
layernormalization107: LayerNorm,
linear193: Linear,
linear194: Linear,
linear195: Linear,
linear196: Linear,
constant528: burn::module::Param<Tensor<2>>,
constant529: burn::module::Param<Tensor<2>>,
constant2592: burn::module::Param<Tensor<1>>,
linear197: Linear,
#[module(skip)]
device: Device,
}
impl Submodule26 {
#[allow(unused_variables)]
pub fn new(device: &Device) -> Self {
let layernormalization103 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let conv1d61 = Conv1dConfig::new(1024, 2048, 1)
.with_stride(1)
.with_padding(PaddingConfig1d::Valid)
.with_dilation(1)
.with_groups(1)
.with_bias(true)
.init(device);
let conv1d62 = Conv1dConfig::new(1024, 1024, 9)
.with_stride(1)
.with_padding(PaddingConfig1d::Valid)
.with_dilation(1)
.with_groups(1024)
.with_bias(true)
.init(device);
let conv1d63 = Conv1dConfig::new(1024, 1024, 1)
.with_stride(1)
.with_padding(PaddingConfig1d::Valid)
.with_dilation(1)
.with_groups(1)
.with_bias(true)
.init(device);
let layernormalization104 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let linear189 = LinearConfig::new(1024, 4096).with_bias(true).init(device);
let linear190 = LinearConfig::new(4096, 1024).with_bias(true).init(device);
let constant2545: 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([0.5f64]),
(device, burn::tensor::DType::F32),
),
device.clone(),
false,
[1].into(),
);
let layernormalization105 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let layernormalization106 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let linear191 = LinearConfig::new(1024, 4096).with_bias(true).init(device);
let linear192 = LinearConfig::new(4096, 1024).with_bias(true).init(device);
let constant2546: 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([0.5f64]),
(device, burn::tensor::DType::F32),
),
device.clone(),
false,
[1].into(),
);
let layernormalization107 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let linear193 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
let linear194 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
let linear195 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
let linear196 = LinearConfig::new(1024, 1024).with_bias(false).init(device);
let constant528: burn::module::Param<Tensor<2>> = burn::module::Param::uninitialized(
burn::module::ParamId::new(),
move |device, _require_grad| Tensor::<
2,
>::zeros([8, 128], (device, burn::tensor::DType::F32)),
device.clone(),
false,
[8, 128].into(),
);
let constant529: burn::module::Param<Tensor<2>> = burn::module::Param::uninitialized(
burn::module::ParamId::new(),
move |device, _require_grad| Tensor::<
2,
>::zeros([8, 128], (device, burn::tensor::DType::F32)),
device.clone(),
false,
[8, 128].into(),
);
let constant2592: 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([11.313708305358887f64]),
(device, burn::tensor::DType::F32),
),
device.clone(),
false,
[1].into(),
);
let linear197 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
Self {
layernormalization103,
conv1d61,
conv1d62,
conv1d63,
layernormalization104,
linear189,
linear190,
constant2545,
layernormalization105,
layernormalization106,
linear191,
linear192,
constant2546,
layernormalization107,
linear193,
linear194,
linear195,
linear196,
constant528,
constant529,
constant2592,
linear197,
device: device.clone(),
}
}
#[allow(clippy::let_and_return, clippy::approx_constant)]
pub fn forward(
&self,
add155_out1: Tensor<3>,
unsqueeze120_out1: Tensor<3, Bool>,
slice1_out1: Tensor<3>,
unsqueeze118_out1: Tensor<4, Bool>,
) -> Tensor<3> {
let layernormalization103_out1 = {
let dtype = add155_out1.clone().dtype();
self.layernormalization103
.forward(add155_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let transpose211_out1 = layernormalization103_out1.permute([0, 2, 1]);
let conv1d61_out1 = self.conv1d61.forward(transpose211_out1);
let split_tensors = conv1d61_out1.split_with_sizes([1024, 1024].into(), 1);
let [split21_out1, split21_out2] = split_tensors.try_into().unwrap();
let sigmoid82_out1 = burn::tensor::activation::sigmoid(split21_out2);
let mul169_out1 = split21_out1.mul(sigmoid82_out1);
let constant2536_out1 = 0f32;
let where81_out1 = mul169_out1.mask_fill(unsqueeze120_out1, constant2536_out1);
let pad42_out1 = where81_out1
.pad(
[(0usize, 0usize), (0usize, 0usize), (4usize, 4usize)],
burn::tensor::ops::PadMode::Constant(0f32),
);
let conv1d62_out1 = self.conv1d62.forward(pad42_out1);
let sigmoid83_out1 = burn::tensor::activation::sigmoid(conv1d62_out1.clone());
let mul170_out1 = conv1d62_out1.mul(sigmoid83_out1);
let conv1d63_out1 = self.conv1d63.forward(mul170_out1);
let transpose213_out1 = conv1d63_out1.permute([0, 2, 1]);
let add156_out1 = add155_out1.add(transpose213_out1);
let layernormalization104_out1 = {
let dtype = add156_out1.clone().dtype();
self.layernormalization104
.forward(add156_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let linear189_out1 = self.linear189.forward(layernormalization104_out1);
let sigmoid84_out1 = burn::tensor::activation::sigmoid(linear189_out1.clone());
let mul171_out1 = linear189_out1.mul(sigmoid84_out1);
let linear190_out1 = self.linear190.forward(mul171_out1);
let constant2545_out1 = self.constant2545.val();
let mul172_out1 = linear190_out1
.mul((constant2545_out1).unsqueeze_dims(&[0isize, 1isize]));
let add157_out1 = add156_out1.add(mul172_out1);
let layernormalization105_out1 = {
let dtype = add157_out1.dtype();
self.layernormalization105
.forward(add157_out1.cast(burn::tensor::DType::F32))
.cast(dtype)
};
let layernormalization106_out1 = {
let dtype = layernormalization105_out1.clone().dtype();
self.layernormalization106
.forward(
layernormalization105_out1.clone().cast(burn::tensor::DType::F32),
)
.cast(dtype)
};
let linear191_out1 = self.linear191.forward(layernormalization106_out1);
let sigmoid85_out1 = burn::tensor::activation::sigmoid(linear191_out1.clone());
let mul173_out1 = linear191_out1.mul(sigmoid85_out1);
let linear192_out1 = self.linear192.forward(mul173_out1);
let constant2546_out1 = self.constant2546.val();
let mul174_out1 = linear192_out1
.mul((constant2546_out1).unsqueeze_dims(&[0isize, 1isize]));
let add158_out1 = layernormalization105_out1.add(mul174_out1);
let layernormalization107_out1 = {
let dtype = add158_out1.clone().dtype();
self.layernormalization107
.forward(add158_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let shape209_out1: [i64; 3] = {
let axes = &layernormalization107_out1.clone().dims()[0..3];
let mut output = [0i64; 3];
for i in 0..3 {
output[i] = axes[i] as i64;
}
output
};
let gather169_out1 = shape209_out1[0] as i64;
let linear193_out1 = self.linear193.forward(layernormalization107_out1.clone());
let unsqueeze361_out1 = [gather169_out1 as i64];
let constant2551_out1: [i64; 1] = [128i64];
let constant2549_out1: [i64; 1] = [-1i64];
let constant2550_out1: [i64; 1] = [8i64];
let concat215_out1: [i64; 4usize] = [
&unsqueeze361_out1[..],
&constant2549_out1[..],
&constant2550_out1[..],
&constant2551_out1[..],
]
.concat()
.try_into()
.unwrap();
let unsqueeze362_out1 = [gather169_out1 as i64];
let constant2555_out1: [i64; 1] = [128i64];
let constant2553_out1: [i64; 1] = [-1i64];
let constant2554_out1: [i64; 1] = [8i64];
let concat216_out1: [i64; 4usize] = [
&unsqueeze362_out1[..],
&constant2553_out1[..],
&constant2554_out1[..],
&constant2555_out1[..],
]
.concat()
.try_into()
.unwrap();
let unsqueeze363_out1 = [gather169_out1 as i64];
let constant2559_out1: [i64; 1] = [128i64];
let constant2557_out1: [i64; 1] = [-1i64];
let constant2558_out1: [i64; 1] = [8i64];
let concat217_out1: [i64; 4usize] = [
&unsqueeze363_out1[..],
&constant2557_out1[..],
&constant2558_out1[..],
&constant2559_out1[..],
]
.concat()
.try_into()
.unwrap();
let reshape251_out1 = linear193_out1.reshape(concat215_out1);
let linear194_out1 = self.linear194.forward(layernormalization107_out1.clone());
let reshape252_out1 = linear194_out1.reshape(concat216_out1);
let linear195_out1 = self.linear195.forward(layernormalization107_out1);
let reshape253_out1 = linear195_out1.reshape(concat217_out1);
let transpose214_out1 = reshape253_out1.permute([0, 2, 1, 3]);
let linear196_out1 = self.linear196.forward(slice1_out1);
let reshape254_out1 = linear196_out1.reshape([1, -1, 8, 128]);
let constant528_out1 = self.constant528.val();
let add159_out1 = reshape251_out1
.clone()
.add((constant528_out1).unsqueeze_dims(&[0isize, 1isize]));
let transpose215_out1 = add159_out1.permute([0, 2, 1, 3]);
let constant529_out1 = self.constant529.val();
let add160_out1 = reshape251_out1
.add((constant529_out1).unsqueeze_dims(&[0isize, 1isize]));
let transpose216_out1 = add160_out1.permute([0, 2, 1, 3]);
let transpose217_out1 = reshape254_out1.permute([0, 2, 3, 1]);
let matmul260_out1 = transpose216_out1.matmul(transpose217_out1);
let shape210_out1: [i64; 4] = {
let axes = &matmul260_out1.clone().dims()[0..4];
let mut output = [0i64; 4];
for i in 0..4 {
output[i] = axes[i] as i64;
}
output
};
let gather170_out1 = shape210_out1[0] as i64;
let gather171_out1 = shape210_out1[1] as i64;
let gather172_out1 = shape210_out1[2] as i64;
let gather173_out1 = shape210_out1[3] as i64;
let pad43_out1 = matmul260_out1
.pad(
[(0usize, 0usize), (0usize, 0usize), (0usize, 0usize), (1usize, 0usize)],
burn::tensor::ops::PadMode::Constant(0f32),
);
let unsqueeze364_out1 = [gather170_out1 as i64];
let unsqueeze365_out1 = [gather171_out1 as i64];
let unsqueeze366_out1 = [gather172_out1 as i64];
let constant2574_out1: [i64; 1] = [-1i64];
let concat219_out1: [i64; 4usize] = [
&unsqueeze364_out1[..],
&unsqueeze365_out1[..],
&constant2574_out1[..],
&unsqueeze366_out1[..],
]
.concat()
.try_into()
.unwrap();
let reshape257_out1 = pad43_out1.reshape(concat219_out1);
let slice110_out1 = reshape257_out1.slice(s![.., .., 1.., ..]);
let unsqueeze367_out1 = [gather170_out1 as i64];
let unsqueeze368_out1 = [gather171_out1 as i64];
let unsqueeze369_out1 = [gather172_out1 as i64];
let unsqueeze370_out1 = [gather173_out1 as i64];
let concat220_out1: [i64; 4usize] = [
&unsqueeze367_out1[..],
&unsqueeze368_out1[..],
&unsqueeze369_out1[..],
&unsqueeze370_out1[..],
]
.concat()
.try_into()
.unwrap();
let reshape258_out1 = slice110_out1.reshape(concat220_out1);
let transpose219_out1 = reshape252_out1.permute([0, 2, 3, 1]);
let matmul261_out1 = transpose215_out1.matmul(transpose219_out1);
let shape214_out1: [i64; 4] = {
let axes = &matmul261_out1.clone().dims()[0..4];
let mut output = [0i64; 4];
for i in 0..4 {
output[i] = axes[i] as i64;
}
output
};
let slice111_out1: [i64; 1] = shape214_out1[3..4].try_into().unwrap();
let squeeze22_out1 = slice111_out1[0] as i64;
let unsqueeze371_out1 = [squeeze22_out1 as i64];
let slice112_out1 = reshape258_out1
.slice(s![.., .., .., 0..unsqueeze371_out1[0]]);
let add161_out1 = matmul261_out1.add(slice112_out1);
let constant2592_out1 = self.constant2592.val();
let div25_out1 = add161_out1
.div((constant2592_out1).unsqueeze_dims(&[0isize, 1isize, 2isize]));
let shape215_out1: [i64; 4] = {
let axes = &transpose214_out1.clone().dims()[0..4];
let mut output = [0i64; 4];
for i in 0..4 {
output[i] = axes[i] as i64;
}
output
};
let gather174_out1 = shape215_out1[0] as i64;
let constant2594_out1 = -10000f32;
let where82_out1 = div25_out1
.mask_fill(unsqueeze118_out1.clone(), constant2594_out1);
let softmax22_out1 = burn::tensor::activation::softmax(where82_out1, 3);
let constant2595_out1 = 0f32;
let where83_out1 = softmax22_out1
.mask_fill(unsqueeze118_out1, constant2595_out1);
let matmul262_out1 = where83_out1.matmul(transpose214_out1);
let transpose220_out1 = matmul262_out1.permute([0, 2, 1, 3]);
let unsqueeze372_out1 = [gather174_out1 as i64];
let constant2598_out1: [i64; 1] = [1024i64];
let constant2597_out1: [i64; 1] = [-1i64];
let concat221_out1: [i64; 3usize] = [
&unsqueeze372_out1[..],
&constant2597_out1[..],
&constant2598_out1[..],
]
.concat()
.try_into()
.unwrap();
let reshape259_out1 = transpose220_out1.reshape(concat221_out1);
let linear197_out1 = self.linear197.forward(reshape259_out1);
let add162_out1 = add158_out1.add(linear197_out1);
add162_out1
}
}
#[derive(Module, Debug)]
pub struct Submodule27 {
layernormalization108: LayerNorm,
conv1d64: Conv1d,
conv1d65: Conv1d,
conv1d66: Conv1d,
layernormalization109: LayerNorm,
linear198: Linear,
linear199: Linear,
constant2608: burn::module::Param<Tensor<1>>,
layernormalization110: LayerNorm,
layernormalization111: LayerNorm,
linear200: Linear,
linear201: Linear,
constant2609: burn::module::Param<Tensor<1>>,
layernormalization112: LayerNorm,
linear202: Linear,
linear203: Linear,
linear204: Linear,
linear205: Linear,
constant552: burn::module::Param<Tensor<2>>,
constant553: burn::module::Param<Tensor<2>>,
constant2655: burn::module::Param<Tensor<1>>,
linear206: Linear,
#[module(skip)]
device: Device,
}
impl Submodule27 {
#[allow(unused_variables)]
pub fn new(device: &Device) -> Self {
let layernormalization108 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let conv1d64 = Conv1dConfig::new(1024, 2048, 1)
.with_stride(1)
.with_padding(PaddingConfig1d::Valid)
.with_dilation(1)
.with_groups(1)
.with_bias(true)
.init(device);
let conv1d65 = Conv1dConfig::new(1024, 1024, 9)
.with_stride(1)
.with_padding(PaddingConfig1d::Valid)
.with_dilation(1)
.with_groups(1024)
.with_bias(true)
.init(device);
let conv1d66 = Conv1dConfig::new(1024, 1024, 1)
.with_stride(1)
.with_padding(PaddingConfig1d::Valid)
.with_dilation(1)
.with_groups(1)
.with_bias(true)
.init(device);
let layernormalization109 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let linear198 = LinearConfig::new(1024, 4096).with_bias(true).init(device);
let linear199 = LinearConfig::new(4096, 1024).with_bias(true).init(device);
let constant2608: 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([0.5f64]),
(device, burn::tensor::DType::F32),
),
device.clone(),
false,
[1].into(),
);
let layernormalization110 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let layernormalization111 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let linear200 = LinearConfig::new(1024, 4096).with_bias(true).init(device);
let linear201 = LinearConfig::new(4096, 1024).with_bias(true).init(device);
let constant2609: 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([0.5f64]),
(device, burn::tensor::DType::F32),
),
device.clone(),
false,
[1].into(),
);
let layernormalization112 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let linear202 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
let linear203 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
let linear204 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
let linear205 = LinearConfig::new(1024, 1024).with_bias(false).init(device);
let constant552: burn::module::Param<Tensor<2>> = burn::module::Param::uninitialized(
burn::module::ParamId::new(),
move |device, _require_grad| Tensor::<
2,
>::zeros([8, 128], (device, burn::tensor::DType::F32)),
device.clone(),
false,
[8, 128].into(),
);
let constant553: burn::module::Param<Tensor<2>> = burn::module::Param::uninitialized(
burn::module::ParamId::new(),
move |device, _require_grad| Tensor::<
2,
>::zeros([8, 128], (device, burn::tensor::DType::F32)),
device.clone(),
false,
[8, 128].into(),
);
let constant2655: 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([11.313708305358887f64]),
(device, burn::tensor::DType::F32),
),
device.clone(),
false,
[1].into(),
);
let linear206 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
Self {
layernormalization108,
conv1d64,
conv1d65,
conv1d66,
layernormalization109,
linear198,
linear199,
constant2608,
layernormalization110,
layernormalization111,
linear200,
linear201,
constant2609,
layernormalization112,
linear202,
linear203,
linear204,
linear205,
constant552,
constant553,
constant2655,
linear206,
device: device.clone(),
}
}
#[allow(clippy::let_and_return, clippy::approx_constant)]
pub fn forward(
&self,
add162_out1: Tensor<3>,
unsqueeze120_out1: Tensor<3, Bool>,
slice1_out1: Tensor<3>,
unsqueeze118_out1: Tensor<4, Bool>,
) -> Tensor<3> {
let layernormalization108_out1 = {
let dtype = add162_out1.clone().dtype();
self.layernormalization108
.forward(add162_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let transpose221_out1 = layernormalization108_out1.permute([0, 2, 1]);
let conv1d64_out1 = self.conv1d64.forward(transpose221_out1);
let split_tensors = conv1d64_out1.split_with_sizes([1024, 1024].into(), 1);
let [split22_out1, split22_out2] = split_tensors.try_into().unwrap();
let sigmoid86_out1 = burn::tensor::activation::sigmoid(split22_out2);
let mul175_out1 = split22_out1.mul(sigmoid86_out1);
let constant2599_out1 = 0f32;
let where84_out1 = mul175_out1.mask_fill(unsqueeze120_out1, constant2599_out1);
let pad44_out1 = where84_out1
.pad(
[(0usize, 0usize), (0usize, 0usize), (4usize, 4usize)],
burn::tensor::ops::PadMode::Constant(0f32),
);
let conv1d65_out1 = self.conv1d65.forward(pad44_out1);
let sigmoid87_out1 = burn::tensor::activation::sigmoid(conv1d65_out1.clone());
let mul176_out1 = conv1d65_out1.mul(sigmoid87_out1);
let conv1d66_out1 = self.conv1d66.forward(mul176_out1);
let transpose223_out1 = conv1d66_out1.permute([0, 2, 1]);
let add163_out1 = add162_out1.add(transpose223_out1);
let layernormalization109_out1 = {
let dtype = add163_out1.clone().dtype();
self.layernormalization109
.forward(add163_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let linear198_out1 = self.linear198.forward(layernormalization109_out1);
let sigmoid88_out1 = burn::tensor::activation::sigmoid(linear198_out1.clone());
let mul177_out1 = linear198_out1.mul(sigmoid88_out1);
let linear199_out1 = self.linear199.forward(mul177_out1);
let constant2608_out1 = self.constant2608.val();
let mul178_out1 = linear199_out1
.mul((constant2608_out1).unsqueeze_dims(&[0isize, 1isize]));
let add164_out1 = add163_out1.add(mul178_out1);
let layernormalization110_out1 = {
let dtype = add164_out1.dtype();
self.layernormalization110
.forward(add164_out1.cast(burn::tensor::DType::F32))
.cast(dtype)
};
let layernormalization111_out1 = {
let dtype = layernormalization110_out1.clone().dtype();
self.layernormalization111
.forward(
layernormalization110_out1.clone().cast(burn::tensor::DType::F32),
)
.cast(dtype)
};
let linear200_out1 = self.linear200.forward(layernormalization111_out1);
let sigmoid89_out1 = burn::tensor::activation::sigmoid(linear200_out1.clone());
let mul179_out1 = linear200_out1.mul(sigmoid89_out1);
let linear201_out1 = self.linear201.forward(mul179_out1);
let constant2609_out1 = self.constant2609.val();
let mul180_out1 = linear201_out1
.mul((constant2609_out1).unsqueeze_dims(&[0isize, 1isize]));
let add165_out1 = layernormalization110_out1.add(mul180_out1);
let layernormalization112_out1 = {
let dtype = add165_out1.clone().dtype();
self.layernormalization112
.forward(add165_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let shape216_out1: [i64; 3] = {
let axes = &layernormalization112_out1.clone().dims()[0..3];
let mut output = [0i64; 3];
for i in 0..3 {
output[i] = axes[i] as i64;
}
output
};
let gather175_out1 = shape216_out1[0] as i64;
let linear202_out1 = self.linear202.forward(layernormalization112_out1.clone());
let unsqueeze373_out1 = [gather175_out1 as i64];
let constant2614_out1: [i64; 1] = [128i64];
let constant2612_out1: [i64; 1] = [-1i64];
let constant2613_out1: [i64; 1] = [8i64];
let concat223_out1: [i64; 4usize] = [
&unsqueeze373_out1[..],
&constant2612_out1[..],
&constant2613_out1[..],
&constant2614_out1[..],
]
.concat()
.try_into()
.unwrap();
let unsqueeze374_out1 = [gather175_out1 as i64];
let constant2618_out1: [i64; 1] = [128i64];
let constant2616_out1: [i64; 1] = [-1i64];
let constant2617_out1: [i64; 1] = [8i64];
let concat224_out1: [i64; 4usize] = [
&unsqueeze374_out1[..],
&constant2616_out1[..],
&constant2617_out1[..],
&constant2618_out1[..],
]
.concat()
.try_into()
.unwrap();
let unsqueeze375_out1 = [gather175_out1 as i64];
let constant2622_out1: [i64; 1] = [128i64];
let constant2620_out1: [i64; 1] = [-1i64];
let constant2621_out1: [i64; 1] = [8i64];
let concat225_out1: [i64; 4usize] = [
&unsqueeze375_out1[..],
&constant2620_out1[..],
&constant2621_out1[..],
&constant2622_out1[..],
]
.concat()
.try_into()
.unwrap();
let reshape262_out1 = linear202_out1.reshape(concat223_out1);
let linear203_out1 = self.linear203.forward(layernormalization112_out1.clone());
let reshape263_out1 = linear203_out1.reshape(concat224_out1);
let linear204_out1 = self.linear204.forward(layernormalization112_out1);
let reshape264_out1 = linear204_out1.reshape(concat225_out1);
let transpose224_out1 = reshape264_out1.permute([0, 2, 1, 3]);
let linear205_out1 = self.linear205.forward(slice1_out1);
let reshape265_out1 = linear205_out1.reshape([1, -1, 8, 128]);
let constant552_out1 = self.constant552.val();
let add166_out1 = reshape262_out1
.clone()
.add((constant552_out1).unsqueeze_dims(&[0isize, 1isize]));
let transpose225_out1 = add166_out1.permute([0, 2, 1, 3]);
let constant553_out1 = self.constant553.val();
let add167_out1 = reshape262_out1
.add((constant553_out1).unsqueeze_dims(&[0isize, 1isize]));
let transpose226_out1 = add167_out1.permute([0, 2, 1, 3]);
let transpose227_out1 = reshape265_out1.permute([0, 2, 3, 1]);
let matmul272_out1 = transpose226_out1.matmul(transpose227_out1);
let shape217_out1: [i64; 4] = {
let axes = &matmul272_out1.clone().dims()[0..4];
let mut output = [0i64; 4];
for i in 0..4 {
output[i] = axes[i] as i64;
}
output
};
let gather176_out1 = shape217_out1[0] as i64;
let gather177_out1 = shape217_out1[1] as i64;
let gather178_out1 = shape217_out1[2] as i64;
let gather179_out1 = shape217_out1[3] as i64;
let pad45_out1 = matmul272_out1
.pad(
[(0usize, 0usize), (0usize, 0usize), (0usize, 0usize), (1usize, 0usize)],
burn::tensor::ops::PadMode::Constant(0f32),
);
let unsqueeze376_out1 = [gather176_out1 as i64];
let unsqueeze377_out1 = [gather177_out1 as i64];
let unsqueeze378_out1 = [gather178_out1 as i64];
let constant2637_out1: [i64; 1] = [-1i64];
let concat227_out1: [i64; 4usize] = [
&unsqueeze376_out1[..],
&unsqueeze377_out1[..],
&constant2637_out1[..],
&unsqueeze378_out1[..],
]
.concat()
.try_into()
.unwrap();
let reshape268_out1 = pad45_out1.reshape(concat227_out1);
let slice115_out1 = reshape268_out1.slice(s![.., .., 1.., ..]);
let unsqueeze379_out1 = [gather176_out1 as i64];
let unsqueeze380_out1 = [gather177_out1 as i64];
let unsqueeze381_out1 = [gather178_out1 as i64];
let unsqueeze382_out1 = [gather179_out1 as i64];
let concat228_out1: [i64; 4usize] = [
&unsqueeze379_out1[..],
&unsqueeze380_out1[..],
&unsqueeze381_out1[..],
&unsqueeze382_out1[..],
]
.concat()
.try_into()
.unwrap();
let reshape269_out1 = slice115_out1.reshape(concat228_out1);
let transpose229_out1 = reshape263_out1.permute([0, 2, 3, 1]);
let matmul273_out1 = transpose225_out1.matmul(transpose229_out1);
let shape221_out1: [i64; 4] = {
let axes = &matmul273_out1.clone().dims()[0..4];
let mut output = [0i64; 4];
for i in 0..4 {
output[i] = axes[i] as i64;
}
output
};
let slice116_out1: [i64; 1] = shape221_out1[3..4].try_into().unwrap();
let squeeze23_out1 = slice116_out1[0] as i64;
let unsqueeze383_out1 = [squeeze23_out1 as i64];
let slice117_out1 = reshape269_out1
.slice(s![.., .., .., 0..unsqueeze383_out1[0]]);
let add168_out1 = matmul273_out1.add(slice117_out1);
let constant2655_out1 = self.constant2655.val();
let div26_out1 = add168_out1
.div((constant2655_out1).unsqueeze_dims(&[0isize, 1isize, 2isize]));
let shape222_out1: [i64; 4] = {
let axes = &transpose224_out1.clone().dims()[0..4];
let mut output = [0i64; 4];
for i in 0..4 {
output[i] = axes[i] as i64;
}
output
};
let gather180_out1 = shape222_out1[0] as i64;
let constant2657_out1 = -10000f32;
let where85_out1 = div26_out1
.mask_fill(unsqueeze118_out1.clone(), constant2657_out1);
let softmax23_out1 = burn::tensor::activation::softmax(where85_out1, 3);
let constant2658_out1 = 0f32;
let where86_out1 = softmax23_out1
.mask_fill(unsqueeze118_out1, constant2658_out1);
let matmul274_out1 = where86_out1.matmul(transpose224_out1);
let transpose230_out1 = matmul274_out1.permute([0, 2, 1, 3]);
let unsqueeze384_out1 = [gather180_out1 as i64];
let constant2661_out1: [i64; 1] = [1024i64];
let constant2660_out1: [i64; 1] = [-1i64];
let concat229_out1: [i64; 3usize] = [
&unsqueeze384_out1[..],
&constant2660_out1[..],
&constant2661_out1[..],
]
.concat()
.try_into()
.unwrap();
let reshape270_out1 = transpose230_out1.reshape(concat229_out1);
let linear206_out1 = self.linear206.forward(reshape270_out1);
let add169_out1 = add165_out1.add(linear206_out1);
add169_out1
}
}
#[derive(Module, Debug)]
pub struct Submodule28 {
layernormalization113: LayerNorm,
conv1d67: Conv1d,
conv1d68: Conv1d,
conv1d69: Conv1d,
layernormalization114: LayerNorm,
linear207: Linear,
linear208: Linear,
constant2671: burn::module::Param<Tensor<1>>,
layernormalization115: LayerNorm,
layernormalization116: LayerNorm,
linear209: Linear,
linear210: Linear,
constant2672: burn::module::Param<Tensor<1>>,
layernormalization117: LayerNorm,
linear211: Linear,
linear212: Linear,
linear213: Linear,
linear214: Linear,
constant576: burn::module::Param<Tensor<2>>,
constant577: burn::module::Param<Tensor<2>>,
#[module(skip)]
device: Device,
}
impl Submodule28 {
#[allow(unused_variables)]
pub fn new(device: &Device) -> Self {
let layernormalization113 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let conv1d67 = Conv1dConfig::new(1024, 2048, 1)
.with_stride(1)
.with_padding(PaddingConfig1d::Valid)
.with_dilation(1)
.with_groups(1)
.with_bias(true)
.init(device);
let conv1d68 = Conv1dConfig::new(1024, 1024, 9)
.with_stride(1)
.with_padding(PaddingConfig1d::Valid)
.with_dilation(1)
.with_groups(1024)
.with_bias(true)
.init(device);
let conv1d69 = Conv1dConfig::new(1024, 1024, 1)
.with_stride(1)
.with_padding(PaddingConfig1d::Valid)
.with_dilation(1)
.with_groups(1)
.with_bias(true)
.init(device);
let layernormalization114 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let linear207 = LinearConfig::new(1024, 4096).with_bias(true).init(device);
let linear208 = LinearConfig::new(4096, 1024).with_bias(true).init(device);
let constant2671: 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([0.5f64]),
(device, burn::tensor::DType::F32),
),
device.clone(),
false,
[1].into(),
);
let layernormalization115 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let layernormalization116 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let linear209 = LinearConfig::new(1024, 4096).with_bias(true).init(device);
let linear210 = LinearConfig::new(4096, 1024).with_bias(true).init(device);
let constant2672: 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([0.5f64]),
(device, burn::tensor::DType::F32),
),
device.clone(),
false,
[1].into(),
);
let layernormalization117 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let linear211 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
let linear212 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
let linear213 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
let linear214 = LinearConfig::new(1024, 1024).with_bias(false).init(device);
let constant576: burn::module::Param<Tensor<2>> = burn::module::Param::uninitialized(
burn::module::ParamId::new(),
move |device, _require_grad| Tensor::<
2,
>::zeros([8, 128], (device, burn::tensor::DType::F32)),
device.clone(),
false,
[8, 128].into(),
);
let constant577: burn::module::Param<Tensor<2>> = burn::module::Param::uninitialized(
burn::module::ParamId::new(),
move |device, _require_grad| Tensor::<
2,
>::zeros([8, 128], (device, burn::tensor::DType::F32)),
device.clone(),
false,
[8, 128].into(),
);
Self {
layernormalization113,
conv1d67,
conv1d68,
conv1d69,
layernormalization114,
linear207,
linear208,
constant2671,
layernormalization115,
layernormalization116,
linear209,
linear210,
constant2672,
layernormalization117,
linear211,
linear212,
linear213,
linear214,
constant576,
constant577,
device: device.clone(),
}
}
#[allow(clippy::let_and_return, clippy::approx_constant)]
pub fn forward(
&self,
add169_out1: Tensor<3>,
unsqueeze120_out1: Tensor<3, Bool>,
slice1_out1: Tensor<3>,
) -> (Tensor<4>, Tensor<4>, Tensor<4>, Tensor<4>, Tensor<3>) {
let layernormalization113_out1 = {
let dtype = add169_out1.clone().dtype();
self.layernormalization113
.forward(add169_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let transpose231_out1 = layernormalization113_out1.permute([0, 2, 1]);
let conv1d67_out1 = self.conv1d67.forward(transpose231_out1);
let split_tensors = conv1d67_out1.split_with_sizes([1024, 1024].into(), 1);
let [split23_out1, split23_out2] = split_tensors.try_into().unwrap();
let sigmoid90_out1 = burn::tensor::activation::sigmoid(split23_out2);
let mul181_out1 = split23_out1.mul(sigmoid90_out1);
let constant2662_out1 = 0f32;
let where87_out1 = mul181_out1.mask_fill(unsqueeze120_out1, constant2662_out1);
let pad46_out1 = where87_out1
.pad(
[(0usize, 0usize), (0usize, 0usize), (4usize, 4usize)],
burn::tensor::ops::PadMode::Constant(0f32),
);
let conv1d68_out1 = self.conv1d68.forward(pad46_out1);
let sigmoid91_out1 = burn::tensor::activation::sigmoid(conv1d68_out1.clone());
let mul182_out1 = conv1d68_out1.mul(sigmoid91_out1);
let conv1d69_out1 = self.conv1d69.forward(mul182_out1);
let transpose233_out1 = conv1d69_out1.permute([0, 2, 1]);
let add170_out1 = add169_out1.add(transpose233_out1);
let layernormalization114_out1 = {
let dtype = add170_out1.clone().dtype();
self.layernormalization114
.forward(add170_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let linear207_out1 = self.linear207.forward(layernormalization114_out1);
let sigmoid92_out1 = burn::tensor::activation::sigmoid(linear207_out1.clone());
let mul183_out1 = linear207_out1.mul(sigmoid92_out1);
let linear208_out1 = self.linear208.forward(mul183_out1);
let constant2671_out1 = self.constant2671.val();
let mul184_out1 = linear208_out1
.mul((constant2671_out1).unsqueeze_dims(&[0isize, 1isize]));
let add171_out1 = add170_out1.add(mul184_out1);
let layernormalization115_out1 = {
let dtype = add171_out1.dtype();
self.layernormalization115
.forward(add171_out1.cast(burn::tensor::DType::F32))
.cast(dtype)
};
let layernormalization116_out1 = {
let dtype = layernormalization115_out1.clone().dtype();
self.layernormalization116
.forward(
layernormalization115_out1.clone().cast(burn::tensor::DType::F32),
)
.cast(dtype)
};
let linear209_out1 = self.linear209.forward(layernormalization116_out1);
let sigmoid93_out1 = burn::tensor::activation::sigmoid(linear209_out1.clone());
let mul185_out1 = linear209_out1.mul(sigmoid93_out1);
let linear210_out1 = self.linear210.forward(mul185_out1);
let constant2672_out1 = self.constant2672.val();
let mul186_out1 = linear210_out1
.mul((constant2672_out1).unsqueeze_dims(&[0isize, 1isize]));
let add172_out1 = layernormalization115_out1.add(mul186_out1);
let layernormalization117_out1 = {
let dtype = add172_out1.clone().dtype();
self.layernormalization117
.forward(add172_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let shape223_out1: [i64; 3] = {
let axes = &layernormalization117_out1.clone().dims()[0..3];
let mut output = [0i64; 3];
for i in 0..3 {
output[i] = axes[i] as i64;
}
output
};
let gather181_out1 = shape223_out1[0] as i64;
let linear211_out1 = self.linear211.forward(layernormalization117_out1.clone());
let unsqueeze385_out1 = [gather181_out1 as i64];
let constant2677_out1: [i64; 1] = [128i64];
let constant2675_out1: [i64; 1] = [-1i64];
let constant2676_out1: [i64; 1] = [8i64];
let concat231_out1: [i64; 4usize] = [
&unsqueeze385_out1[..],
&constant2675_out1[..],
&constant2676_out1[..],
&constant2677_out1[..],
]
.concat()
.try_into()
.unwrap();
let unsqueeze386_out1 = [gather181_out1 as i64];
let constant2681_out1: [i64; 1] = [128i64];
let constant2679_out1: [i64; 1] = [-1i64];
let constant2680_out1: [i64; 1] = [8i64];
let concat232_out1: [i64; 4usize] = [
&unsqueeze386_out1[..],
&constant2679_out1[..],
&constant2680_out1[..],
&constant2681_out1[..],
]
.concat()
.try_into()
.unwrap();
let unsqueeze387_out1 = [gather181_out1 as i64];
let constant2685_out1: [i64; 1] = [128i64];
let constant2683_out1: [i64; 1] = [-1i64];
let constant2684_out1: [i64; 1] = [8i64];
let concat233_out1: [i64; 4usize] = [
&unsqueeze387_out1[..],
&constant2683_out1[..],
&constant2684_out1[..],
&constant2685_out1[..],
]
.concat()
.try_into()
.unwrap();
let reshape273_out1 = linear211_out1.reshape(concat231_out1);
let linear212_out1 = self.linear212.forward(layernormalization117_out1.clone());
let reshape274_out1 = linear212_out1.reshape(concat232_out1);
let linear213_out1 = self.linear213.forward(layernormalization117_out1);
let reshape275_out1 = linear213_out1.reshape(concat233_out1);
let transpose234_out1 = reshape275_out1.permute([0, 2, 1, 3]);
let linear214_out1 = self.linear214.forward(slice1_out1);
let reshape276_out1 = linear214_out1.reshape([1, -1, 8, 128]);
let constant576_out1 = self.constant576.val();
let add173_out1 = reshape273_out1
.clone()
.add((constant576_out1).unsqueeze_dims(&[0isize, 1isize]));
let transpose235_out1 = add173_out1.permute([0, 2, 1, 3]);
let constant577_out1 = self.constant577.val();
let add174_out1 = reshape273_out1
.add((constant577_out1).unsqueeze_dims(&[0isize, 1isize]));
let transpose236_out1 = add174_out1.permute([0, 2, 1, 3]);
let transpose237_out1 = reshape276_out1.permute([0, 2, 3, 1]);
let matmul284_out1 = transpose236_out1.matmul(transpose237_out1);
(
matmul284_out1,
reshape274_out1,
transpose235_out1,
transpose234_out1,
add172_out1,
)
}
}
#[derive(Module, Debug)]
pub struct Submodule29 {
constant2718: burn::module::Param<Tensor<1>>,
linear215: Linear,
layernormalization118: LayerNorm,
conv1d70: Conv1d,
conv1d71: Conv1d,
conv1d72: Conv1d,
layernormalization119: LayerNorm,
linear216: Linear,
linear217: Linear,
constant2734: burn::module::Param<Tensor<1>>,
layernormalization120: LayerNorm,
#[module(skip)]
device: Device,
}
impl Submodule29 {
#[allow(unused_variables)]
pub fn new(device: &Device) -> Self {
let constant2718: 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([11.313708305358887f64]),
(device, burn::tensor::DType::F32),
),
device.clone(),
false,
[1].into(),
);
let linear215 = LinearConfig::new(1024, 1024).with_bias(true).init(device);
let layernormalization118 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let conv1d70 = Conv1dConfig::new(1024, 2048, 1)
.with_stride(1)
.with_padding(PaddingConfig1d::Valid)
.with_dilation(1)
.with_groups(1)
.with_bias(true)
.init(device);
let conv1d71 = Conv1dConfig::new(1024, 1024, 9)
.with_stride(1)
.with_padding(PaddingConfig1d::Valid)
.with_dilation(1)
.with_groups(1024)
.with_bias(true)
.init(device);
let conv1d72 = Conv1dConfig::new(1024, 1024, 1)
.with_stride(1)
.with_padding(PaddingConfig1d::Valid)
.with_dilation(1)
.with_groups(1)
.with_bias(true)
.init(device);
let layernormalization119 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
let linear216 = LinearConfig::new(1024, 4096).with_bias(true).init(device);
let linear217 = LinearConfig::new(4096, 1024).with_bias(true).init(device);
let constant2734: 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([0.5f64]),
(device, burn::tensor::DType::F32),
),
device.clone(),
false,
[1].into(),
);
let layernormalization120 = LayerNormConfig::new(1024)
.with_epsilon(0.000009999999747378752f64)
.with_bias(true)
.init(device);
Self {
constant2718,
linear215,
layernormalization118,
conv1d70,
conv1d71,
conv1d72,
layernormalization119,
linear216,
linear217,
constant2734,
layernormalization120,
device: device.clone(),
}
}
#[allow(clippy::let_and_return, clippy::approx_constant)]
pub fn forward(
&self,
matmul284_out1: Tensor<4>,
reshape274_out1: Tensor<4>,
transpose235_out1: Tensor<4>,
transpose234_out1: Tensor<4>,
unsqueeze118_out1: Tensor<4, Bool>,
add172_out1: Tensor<3>,
unsqueeze120_out1: Tensor<3, Bool>,
) -> Tensor<3> {
let shape224_out1: [i64; 4] = {
let axes = &matmul284_out1.clone().dims()[0..4];
let mut output = [0i64; 4];
for i in 0..4 {
output[i] = axes[i] as i64;
}
output
};
let gather182_out1 = shape224_out1[0] as i64;
let gather183_out1 = shape224_out1[1] as i64;
let gather184_out1 = shape224_out1[2] as i64;
let gather185_out1 = shape224_out1[3] as i64;
let pad47_out1 = matmul284_out1
.pad(
[(0usize, 0usize), (0usize, 0usize), (0usize, 0usize), (1usize, 0usize)],
burn::tensor::ops::PadMode::Constant(0f32),
);
let unsqueeze388_out1 = [gather182_out1 as i64];
let unsqueeze389_out1 = [gather183_out1 as i64];
let unsqueeze390_out1 = [gather184_out1 as i64];
let constant2700_out1: [i64; 1] = [-1i64];
let concat235_out1: [i64; 4usize] = [
&unsqueeze388_out1[..],
&unsqueeze389_out1[..],
&constant2700_out1[..],
&unsqueeze390_out1[..],
]
.concat()
.try_into()
.unwrap();
let reshape279_out1 = pad47_out1.reshape(concat235_out1);
let slice120_out1 = reshape279_out1.slice(s![.., .., 1.., ..]);
let unsqueeze391_out1 = [gather182_out1 as i64];
let unsqueeze392_out1 = [gather183_out1 as i64];
let unsqueeze393_out1 = [gather184_out1 as i64];
let unsqueeze394_out1 = [gather185_out1 as i64];
let concat236_out1: [i64; 4usize] = [
&unsqueeze391_out1[..],
&unsqueeze392_out1[..],
&unsqueeze393_out1[..],
&unsqueeze394_out1[..],
]
.concat()
.try_into()
.unwrap();
let reshape280_out1 = slice120_out1.reshape(concat236_out1);
let transpose239_out1 = reshape274_out1.permute([0, 2, 3, 1]);
let matmul285_out1 = transpose235_out1.matmul(transpose239_out1);
let shape228_out1: [i64; 4] = {
let axes = &matmul285_out1.clone().dims()[0..4];
let mut output = [0i64; 4];
for i in 0..4 {
output[i] = axes[i] as i64;
}
output
};
let slice121_out1: [i64; 1] = shape228_out1[3..4].try_into().unwrap();
let squeeze24_out1 = slice121_out1[0] as i64;
let unsqueeze395_out1 = [squeeze24_out1 as i64];
let slice122_out1 = reshape280_out1
.slice(s![.., .., .., 0..unsqueeze395_out1[0]]);
let add175_out1 = matmul285_out1.add(slice122_out1);
let constant2718_out1 = self.constant2718.val();
let div27_out1 = add175_out1
.div((constant2718_out1).unsqueeze_dims(&[0isize, 1isize, 2isize]));
let shape229_out1: [i64; 4] = {
let axes = &transpose234_out1.clone().dims()[0..4];
let mut output = [0i64; 4];
for i in 0..4 {
output[i] = axes[i] as i64;
}
output
};
let gather186_out1 = shape229_out1[0] as i64;
let constant2720_out1 = -10000f32;
let where88_out1 = div27_out1
.mask_fill(unsqueeze118_out1.clone(), constant2720_out1);
let softmax24_out1 = burn::tensor::activation::softmax(where88_out1, 3);
let constant2721_out1 = 0f32;
let where89_out1 = softmax24_out1
.mask_fill(unsqueeze118_out1, constant2721_out1);
let matmul286_out1 = where89_out1.matmul(transpose234_out1);
let transpose240_out1 = matmul286_out1.permute([0, 2, 1, 3]);
let unsqueeze396_out1 = [gather186_out1 as i64];
let constant2724_out1: [i64; 1] = [1024i64];
let constant2723_out1: [i64; 1] = [-1i64];
let concat237_out1: [i64; 3usize] = [
&unsqueeze396_out1[..],
&constant2723_out1[..],
&constant2724_out1[..],
]
.concat()
.try_into()
.unwrap();
let reshape281_out1 = transpose240_out1.reshape(concat237_out1);
let linear215_out1 = self.linear215.forward(reshape281_out1);
let add176_out1 = add172_out1.add(linear215_out1);
let layernormalization118_out1 = {
let dtype = add176_out1.clone().dtype();
self.layernormalization118
.forward(add176_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let transpose241_out1 = layernormalization118_out1.permute([0, 2, 1]);
let conv1d70_out1 = self.conv1d70.forward(transpose241_out1);
let split_tensors = conv1d70_out1.split_with_sizes([1024, 1024].into(), 1);
let [split24_out1, split24_out2] = split_tensors.try_into().unwrap();
let sigmoid94_out1 = burn::tensor::activation::sigmoid(split24_out2);
let mul187_out1 = split24_out1.mul(sigmoid94_out1);
let constant2725_out1 = 0f32;
let where90_out1 = mul187_out1.mask_fill(unsqueeze120_out1, constant2725_out1);
let pad48_out1 = where90_out1
.pad(
[(0usize, 0usize), (0usize, 0usize), (4usize, 4usize)],
burn::tensor::ops::PadMode::Constant(0f32),
);
let conv1d71_out1 = self.conv1d71.forward(pad48_out1);
let sigmoid95_out1 = burn::tensor::activation::sigmoid(conv1d71_out1.clone());
let mul188_out1 = conv1d71_out1.mul(sigmoid95_out1);
let conv1d72_out1 = self.conv1d72.forward(mul188_out1);
let transpose243_out1 = conv1d72_out1.permute([0, 2, 1]);
let add177_out1 = add176_out1.add(transpose243_out1);
let layernormalization119_out1 = {
let dtype = add177_out1.clone().dtype();
self.layernormalization119
.forward(add177_out1.clone().cast(burn::tensor::DType::F32))
.cast(dtype)
};
let linear216_out1 = self.linear216.forward(layernormalization119_out1);
let sigmoid96_out1 = burn::tensor::activation::sigmoid(linear216_out1.clone());
let mul189_out1 = linear216_out1.mul(sigmoid96_out1);
let linear217_out1 = self.linear217.forward(mul189_out1);
let constant2734_out1 = self.constant2734.val();
let mul190_out1 = linear217_out1
.mul((constant2734_out1).unsqueeze_dims(&[0isize, 1isize]));
let add178_out1 = add177_out1.add(mul190_out1);
let layernormalization120_out1 = {
let dtype = add178_out1.dtype();
self.layernormalization120
.forward(add178_out1.cast(burn::tensor::DType::F32))
.cast(dtype)
};
let transpose244_out1 = layernormalization120_out1.permute([0, 2, 1]);
transpose244_out1
}
}
#[derive(Module, Debug)]
pub struct Model {
submodule1: Submodule1,
submodule2: Submodule2,
submodule3: Submodule3,
submodule4: Submodule4,
submodule5: Submodule5,
submodule6: Submodule6,
submodule7: Submodule7,
submodule8: Submodule8,
submodule9: Submodule9,
submodule10: Submodule10,
submodule11: Submodule11,
submodule12: Submodule12,
submodule13: Submodule13,
submodule14: Submodule14,
submodule15: Submodule15,
submodule16: Submodule16,
submodule17: Submodule17,
submodule18: Submodule18,
submodule19: Submodule19,
submodule20: Submodule20,
submodule21: Submodule21,
submodule22: Submodule22,
submodule23: Submodule23,
submodule24: Submodule24,
submodule25: Submodule25,
submodule26: Submodule26,
submodule27: Submodule27,
submodule28: Submodule28,
submodule29: Submodule29,
#[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);
let submodule5 = Submodule5::new(device);
let submodule6 = Submodule6::new(device);
let submodule7 = Submodule7::new(device);
let submodule8 = Submodule8::new(device);
let submodule9 = Submodule9::new(device);
let submodule10 = Submodule10::new(device);
let submodule11 = Submodule11::new(device);
let submodule12 = Submodule12::new(device);
let submodule13 = Submodule13::new(device);
let submodule14 = Submodule14::new(device);
let submodule15 = Submodule15::new(device);
let submodule16 = Submodule16::new(device);
let submodule17 = Submodule17::new(device);
let submodule18 = Submodule18::new(device);
let submodule19 = Submodule19::new(device);
let submodule20 = Submodule20::new(device);
let submodule21 = Submodule21::new(device);
let submodule22 = Submodule22::new(device);
let submodule23 = Submodule23::new(device);
let submodule24 = Submodule24::new(device);
let submodule25 = Submodule25::new(device);
let submodule26 = Submodule26::new(device);
let submodule27 = Submodule27::new(device);
let submodule28 = Submodule28::new(device);
let submodule29 = Submodule29::new(device);
Self {
submodule1,
submodule2,
submodule3,
submodule4,
submodule5,
submodule6,
submodule7,
submodule8,
submodule9,
submodule10,
submodule11,
submodule12,
submodule13,
submodule14,
submodule15,
submodule16,
submodule17,
submodule18,
submodule19,
submodule20,
submodule21,
submodule22,
submodule23,
submodule24,
submodule25,
submodule26,
submodule27,
submodule28,
submodule29,
device: device.clone(),
}
}
#[allow(clippy::let_and_return, clippy::approx_constant)]
pub fn forward(
&self,
audio_signal: Tensor<3>,
length: Tensor<1, Int>,
) -> (Tensor<3>, Tensor<1, Int>) {
let (mul17_out1, add3_out1) = self.submodule1.forward(audio_signal, length);
let (mul30_out1, add6_out1) = self.submodule2.forward(mul17_out1, add3_out1);
let (add9_out1, mul43_out1, cast17_out1) = self
.submodule3
.forward(mul30_out1, add6_out1);
let (add11_out1, slice1_out1, not1_out1, not2_out1, cast20_out1) = self
.submodule4
.forward(add9_out1, mul43_out1, cast17_out1);
let (add15_out1, unsqueeze118_out1) = self
.submodule5
.forward(add11_out1, slice1_out1.clone(), not1_out1);
let (add22_out1, unsqueeze120_out1) = self
.submodule6
.forward(
add15_out1,
not2_out1,
slice1_out1.clone(),
unsqueeze118_out1.clone(),
);
let add29_out1 = self
.submodule7
.forward(
add22_out1,
unsqueeze120_out1.clone(),
slice1_out1.clone(),
unsqueeze118_out1.clone(),
);
let add36_out1 = self
.submodule8
.forward(
add29_out1,
unsqueeze120_out1.clone(),
slice1_out1.clone(),
unsqueeze118_out1.clone(),
);
let add43_out1 = self
.submodule9
.forward(
add36_out1,
unsqueeze120_out1.clone(),
slice1_out1.clone(),
unsqueeze118_out1.clone(),
);
let add50_out1 = self
.submodule10
.forward(
add43_out1,
unsqueeze120_out1.clone(),
slice1_out1.clone(),
unsqueeze118_out1.clone(),
);
let add57_out1 = self
.submodule11
.forward(
add50_out1,
unsqueeze120_out1.clone(),
slice1_out1.clone(),
unsqueeze118_out1.clone(),
);
let add64_out1 = self
.submodule12
.forward(
add57_out1,
unsqueeze120_out1.clone(),
slice1_out1.clone(),
unsqueeze118_out1.clone(),
);
let add71_out1 = self
.submodule13
.forward(
add64_out1,
unsqueeze120_out1.clone(),
slice1_out1.clone(),
unsqueeze118_out1.clone(),
);
let add78_out1 = self
.submodule14
.forward(
add71_out1,
unsqueeze120_out1.clone(),
slice1_out1.clone(),
unsqueeze118_out1.clone(),
);
let add85_out1 = self
.submodule15
.forward(
add78_out1,
unsqueeze120_out1.clone(),
slice1_out1.clone(),
unsqueeze118_out1.clone(),
);
let add92_out1 = self
.submodule16
.forward(
add85_out1,
unsqueeze120_out1.clone(),
slice1_out1.clone(),
unsqueeze118_out1.clone(),
);
let add99_out1 = self
.submodule17
.forward(
add92_out1,
unsqueeze120_out1.clone(),
slice1_out1.clone(),
unsqueeze118_out1.clone(),
);
let add106_out1 = self
.submodule18
.forward(
add99_out1,
unsqueeze120_out1.clone(),
slice1_out1.clone(),
unsqueeze118_out1.clone(),
);
let add113_out1 = self
.submodule19
.forward(
add106_out1,
unsqueeze120_out1.clone(),
slice1_out1.clone(),
unsqueeze118_out1.clone(),
);
let add120_out1 = self
.submodule20
.forward(
add113_out1,
unsqueeze120_out1.clone(),
slice1_out1.clone(),
unsqueeze118_out1.clone(),
);
let add127_out1 = self
.submodule21
.forward(
add120_out1,
unsqueeze120_out1.clone(),
slice1_out1.clone(),
unsqueeze118_out1.clone(),
);
let add134_out1 = self
.submodule22
.forward(
add127_out1,
unsqueeze120_out1.clone(),
slice1_out1.clone(),
unsqueeze118_out1.clone(),
);
let add141_out1 = self
.submodule23
.forward(
add134_out1,
unsqueeze120_out1.clone(),
slice1_out1.clone(),
unsqueeze118_out1.clone(),
);
let add148_out1 = self
.submodule24
.forward(
add141_out1,
unsqueeze120_out1.clone(),
slice1_out1.clone(),
unsqueeze118_out1.clone(),
);
let add155_out1 = self
.submodule25
.forward(
add148_out1,
unsqueeze120_out1.clone(),
slice1_out1.clone(),
unsqueeze118_out1.clone(),
);
let add162_out1 = self
.submodule26
.forward(
add155_out1,
unsqueeze120_out1.clone(),
slice1_out1.clone(),
unsqueeze118_out1.clone(),
);
let add169_out1 = self
.submodule27
.forward(
add162_out1,
unsqueeze120_out1.clone(),
slice1_out1.clone(),
unsqueeze118_out1.clone(),
);
let (
matmul284_out1,
reshape274_out1,
transpose235_out1,
transpose234_out1,
add172_out1,
) = self
.submodule28
.forward(add169_out1, unsqueeze120_out1.clone(), slice1_out1);
let transpose244_out1 = self
.submodule29
.forward(
matmul284_out1,
reshape274_out1,
transpose235_out1,
transpose234_out1,
unsqueeze118_out1,
add172_out1,
unsqueeze120_out1,
);
(transpose244_out1, cast20_out1)
}
}