Compile Model Crashing (Transformer)

[mune98]

Hello!

I was trying to compile a custom PyTorch Transformer model on the DeepSparse Engine but kept crashing when compiling the model. I faced the same problem when I used the vanilla Transformer Encoder from PyTorch. Is there no support for Transformer only models? FYI, I tried the same steps for pretrained AlexNet from PyTorch and it worked. My steps are shown below:

Using:
DeepSparse 1.2.0
Torch 1.12.1+cu113
Python 3.7.15

Instantiating model:

encoder_layer = nn.TransformerEncoderLayer(d_model=512, nhead=8)
transformer_encoder = nn.TransformerEncoder(encoder_layer, num_layers=2)
src = torch.rand((1,32,512))
out = transformer_encoder(src)

Convert to ONNX:

input_names = ["Input"]
output_names = ["Output"]

torch.onnx.export(transformer_encoder, 
                  src, 
                  "PyTorchTransformerEncoder.onnx", 
                  verbose=True, 
                  input_names=input_names, 
                  output_names=output_names,
                  export_params=True,
                  dynamic_axes={"Input": {0: 'batch_size'},
                                "Output": {0: 'batch_size'}})

Compile and run on DeepSparse (notebook crashes at compile_model):

from deepsparse import compile_model
from deepsparse.utils import generate_random_inputs

onnx_filepath =  "PyTorchTransformerEncoder.onnx"
batch_size = 16

# Generate random sample input
inputs = generate_random_inputs(onnx_filepath, batch_size)

# # Compile and run
engine = compile_model(onnx_filepath, batch_size)
outputs = engine.run(inputs)

Converted ONNX Graph, using print(onnx.helper.printable_graph(onnx_model.graph)):

graph torch_jit (
  %Input[FLOAT, batch_sizex32x512]
) initializers ( 
...
  return %Output
}

graph torch_jit (
  %Input[FLOAT, batch_sizex32x512]
) initializers (
  %layers.0.self_attn.in_proj_bias[FLOAT, 1536]
  %layers.0.self_attn.out_proj.weight[FLOAT, 512x512]
  %layers.0.self_attn.out_proj.bias[FLOAT, 512]
  %layers.0.linear1.bias[FLOAT, 2048]
  %layers.0.linear2.bias[FLOAT, 512]
  %layers.0.norm1.weight[FLOAT, 512]
  %onnx::MatMul_281[FLOAT, 512x1536]
  %onnx::MatMul_282[FLOAT, 512x2048]
  %onnx::MatMul_283[FLOAT, 2048x512]
) {
  %onnx::MatMul_286 = Identity(%onnx::MatMul_283)
  %onnx::MatMul_285 = Identity(%onnx::MatMul_282)
  %onnx::MatMul_284 = Identity(%onnx::MatMul_281)
  %layers.1.norm2.bias = Identity(%layers.0.self_attn.out_proj.bias)
  %layers.1.norm2.weight = Identity(%layers.0.norm1.weight)
  %layers.1.norm1.bias = Identity(%layers.0.self_attn.out_proj.bias)
  %layers.1.norm1.weight = Identity(%layers.0.norm1.weight)
  %layers.1.linear2.bias = Identity(%layers.0.linear2.bias)
  %layers.1.linear1.bias = Identity(%layers.0.linear1.bias)
  %layers.1.self_attn.out_proj.bias = Identity(%layers.0.self_attn.out_proj.bias)
  %layers.1.self_attn.out_proj.weight = Identity(%layers.0.self_attn.out_proj.weight)
  %layers.1.self_attn.in_proj_bias = Identity(%layers.0.self_attn.in_proj_bias)
  %layers.0.norm2.bias = Identity(%layers.0.self_attn.out_proj.bias)
  %layers.0.norm2.weight = Identity(%layers.0.norm1.weight)
  %layers.0.norm1.bias = Identity(%layers.0.self_attn.out_proj.bias)
  %onnx::Gather_25 = Shape(%Input)
  %onnx::Gather_26 = Constant[value = <Scalar Tensor []>]()
  %onnx::Unsqueeze_27 = Gather[axis = 0](%onnx::Gather_25, %onnx::Gather_26)
  %onnx::Gather_28 = Shape(%Input)
  %onnx::Gather_29 = Constant[value = <Scalar Tensor []>]()
  %onnx::Mul_30 = Gather[axis = 0](%onnx::Gather_28, %onnx::Gather_29)
  %onnx::Gather_31 = Shape(%Input)
  %onnx::Gather_32 = Constant[value = <Scalar Tensor []>]()
  %onnx::Div_33 = Gather[axis = 0](%onnx::Gather_31, %onnx::Gather_32)
  %onnx::Div_34 = Constant[value = <Scalar Tensor []>]()
  %onnx::Cast_35 = Div(%onnx::Div_33, %onnx::Div_34)
  %onnx::Cast_36 = Cast[to = 7](%onnx::Cast_35)
  %onnx::Unsqueeze_37 = Cast[to = 7](%onnx::Cast_36)
  %onnx::Add_39 = MatMul(%Input, %onnx::MatMul_281)
  %onnx::Shape_40 = Add(%layers.0.self_attn.in_proj_bias, %onnx::Add_39)
  %onnx::Gather_41 = Shape(%onnx::Shape_40)
  %onnx::Gather_42 = Constant[value = <Tensor>]()
  %onnx::Add_43 = Gather[axis = 0](%onnx::Gather_41, %onnx::Gather_42)
  %onnx::Slice_44 = Constant[value = <Tensor>]()
  %onnx::Add_45 = Constant[value = <Tensor>]()
  %onnx::Div_46 = Add(%onnx::Add_43, %onnx::Add_45)
  %onnx::Div_47 = Constant[value = <Tensor>]()
  %onnx::Mul_48 = Div(%onnx::Div_46, %onnx::Div_47)
  %onnx::Mul_49 = Constant[value = <Tensor>]()
  %onnx::Slice_50 = Mul(%onnx::Mul_48, %onnx::Mul_49)
  %onnx::Reshape_51 = Slice(%onnx::Shape_40, %onnx::Slice_44, %onnx::Slice_50, %onnx::Gather_42)
  %onnx::Mul_52 = Constant[value = <Tensor>]()
  %onnx::Slice_53 = Mul(%onnx::Mul_48, %onnx::Mul_52)
  %onnx::Shape_54 = Slice(%onnx::Shape_40, %onnx::Slice_50, %onnx::Slice_53, %onnx::Gather_42)
  %onnx::Mul_55 = Constant[value = <Tensor>]()
  %onnx::Slice_56 = Mul(%onnx::Mul_48, %onnx::Mul_55)
  %onnx::Shape_57 = Slice(%onnx::Shape_40, %onnx::Slice_53, %onnx::Slice_56, %onnx::Gather_42)
  %onnx::Mul_58 = Constant[value = <Scalar Tensor []>]()
  %onnx::Unsqueeze_59 = Mul(%onnx::Mul_30, %onnx::Mul_58)
  %onnx::Unsqueeze_60 = Constant[value = <Tensor>]()
  %onnx::Concat_61 = Unsqueeze(%onnx::Unsqueeze_27, %onnx::Unsqueeze_60)
  %onnx::Unsqueeze_62 = Constant[value = <Tensor>]()
  %onnx::Concat_63 = Unsqueeze(%onnx::Unsqueeze_59, %onnx::Unsqueeze_62)
  %onnx::Unsqueeze_64 = Constant[value = <Tensor>]()
  %onnx::Concat_65 = Unsqueeze(%onnx::Unsqueeze_37, %onnx::Unsqueeze_64)
  %onnx::Reshape_66 = Concat[axis = 0](%onnx::Concat_61, %onnx::Concat_63, %onnx::Concat_65)
  %onnx::Transpose_67 = Reshape(%onnx::Reshape_51, %onnx::Reshape_66)
  %q = Transpose[perm = [1, 0, 2]](%onnx::Transpose_67)
  %onnx::Gather_69 = Shape(%onnx::Shape_54)
  %onnx::Gather_70 = Constant[value = <Scalar Tensor []>]()
  %onnx::Unsqueeze_71 = Gather[axis = 0](%onnx::Gather_69, %onnx::Gather_70)
  %onnx::Mul_72 = Constant[value = <Scalar Tensor []>]()
  %onnx::Unsqueeze_73 = Mul(%onnx::Mul_30, %onnx::Mul_72)
  %onnx::Unsqueeze_74 = Constant[value = <Tensor>]()
  %onnx::Concat_75 = Unsqueeze(%onnx::Unsqueeze_71, %onnx::Unsqueeze_74)
  %onnx::Unsqueeze_76 = Constant[value = <Tensor>]()
  %onnx::Concat_77 = Unsqueeze(%onnx::Unsqueeze_73, %onnx::Unsqueeze_76)
  %onnx::Unsqueeze_78 = Constant[value = <Tensor>]()
  %onnx::Concat_79 = Unsqueeze(%onnx::Unsqueeze_37, %onnx::Unsqueeze_78)
  %onnx::Reshape_80 = Concat[axis = 0](%onnx::Concat_75, %onnx::Concat_77, %onnx::Concat_79)
  %onnx::Transpose_81 = Reshape(%onnx::Shape_54, %onnx::Reshape_80)
  %onnx::Gather_82 = Shape(%onnx::Shape_57)
  %onnx::Gather_83 = Constant[value = <Scalar Tensor []>]()
  %onnx::Unsqueeze_84 = Gather[axis = 0](%onnx::Gather_82, %onnx::Gather_83)
  %onnx::Mul_85 = Constant[value = <Scalar Tensor []>]()
  %onnx::Unsqueeze_86 = Mul(%onnx::Mul_30, %onnx::Mul_85)
  %onnx::Unsqueeze_87 = Constant[value = <Tensor>]()
  %onnx::Concat_88 = Unsqueeze(%onnx::Unsqueeze_84, %onnx::Unsqueeze_87)
  %onnx::Unsqueeze_89 = Constant[value = <Tensor>]()
  %onnx::Concat_90 = Unsqueeze(%onnx::Unsqueeze_86, %onnx::Unsqueeze_89)
  %onnx::Unsqueeze_91 = Constant[value = <Tensor>]()
  %onnx::Concat_92 = Unsqueeze(%onnx::Unsqueeze_37, %onnx::Unsqueeze_91)
  %onnx::Reshape_93 = Concat[axis = 0](%onnx::Concat_88, %onnx::Concat_90, %onnx::Concat_92)
  %onnx::Transpose_94 = Reshape(%onnx::Shape_57, %onnx::Reshape_93)
  %v = Transpose[perm = [1, 0, 2]](%onnx::Transpose_94)
  %onnx::Div_96 = Constant[value = <Scalar Tensor []>]()
  %onnx::MatMul_97 = Div(%q, %onnx::Div_96)
  %onnx::MatMul_98 = Transpose[perm = [1, 2, 0]](%onnx::Transpose_81)
  %attn = MatMul(%onnx::MatMul_97, %onnx::MatMul_98)
  %onnx::MatMul_100 = Softmax[axis = -1](%attn)
  %onnx::Transpose_101 = MatMul(%onnx::MatMul_100, %v)
  %onnx::Reshape_102 = Transpose[perm = [1, 0, 2]](%onnx::Transpose_101)
  %onnx::Unsqueeze_103 = Mul(%onnx::Unsqueeze_27, %onnx::Mul_30)
  %onnx::Unsqueeze_104 = Constant[value = <Tensor>]()
  %onnx::Concat_105 = Unsqueeze(%onnx::Unsqueeze_103, %onnx::Unsqueeze_104)
  %onnx::Unsqueeze_106 = Constant[value = <Tensor>]()
  %onnx::Concat_107 = Unsqueeze(%onnx::Div_33, %onnx::Unsqueeze_106)
  %onnx::Reshape_108 = Concat[axis = 0](%onnx::Concat_105, %onnx::Concat_107)
  %onnx::Gemm_109 = Reshape(%onnx::Reshape_102, %onnx::Reshape_108)
  %onnx::Shape_110 = Gemm[alpha = 1, beta = 1, transB = 1](%onnx::Gemm_109, %layers.0.self_attn.out_proj.weight, %layers.0.self_attn.out_proj.bias)
  %onnx::Gather_111 = Shape(%onnx::Shape_110)
  %onnx::Gather_112 = Constant[value = <Scalar Tensor []>]()
  %onnx::Unsqueeze_113 = Gather[axis = 0](%onnx::Gather_111, %onnx::Gather_112)
  %onnx::Unsqueeze_114 = Constant[value = <Tensor>]()
  %onnx::Concat_115 = Unsqueeze(%onnx::Unsqueeze_27, %onnx::Unsqueeze_114)
  %onnx::Unsqueeze_116 = Constant[value = <Tensor>]()
  %onnx::Concat_117 = Unsqueeze(%onnx::Mul_30, %onnx::Unsqueeze_116)
  %onnx::Unsqueeze_118 = Constant[value = <Tensor>]()
  %onnx::Concat_119 = Unsqueeze(%onnx::Unsqueeze_113, %onnx::Unsqueeze_118)
  %onnx::Reshape_120 = Concat[axis = 0](%onnx::Concat_115, %onnx::Concat_117, %onnx::Concat_119)
  %input = Reshape(%onnx::Shape_110, %onnx::Reshape_120)
  %input.3 = Add(%Input, %input)
  %onnx::Sub_123 = ReduceMean[axes = [-1]](%input.3)
  %onnx::Pow_124 = Sub(%input.3, %onnx::Sub_123)
  %onnx::Pow_125 = Constant[value = <Scalar Tensor []>]()
  %onnx::ReduceMean_126 = Pow(%onnx::Pow_124, %onnx::Pow_125)
  %onnx::Add_127 = ReduceMean[axes = [-1]](%onnx::ReduceMean_126)
  %onnx::Add_128 = Constant[value = <Scalar Tensor []>]()
  %onnx::Sqrt_129 = Add(%onnx::Add_127, %onnx::Add_128)
  %onnx::Div_130 = Sqrt(%onnx::Sqrt_129)
  %onnx::Mul_131 = Div(%onnx::Pow_124, %onnx::Div_130)
  %onnx::Add_132 = Mul(%onnx::Mul_131, %layers.0.norm1.weight)
  %onnx::MatMul_133 = Add(%onnx::Add_132, %layers.0.norm1.bias)
  %onnx::Add_135 = MatMul(%onnx::MatMul_133, %onnx::MatMul_282)
  %onnx::Relu_136 = Add(%layers.0.linear1.bias, %onnx::Add_135)
  %input.7 = Relu(%onnx::Relu_136)
  %onnx::Add_139 = MatMul(%input.7, %onnx::MatMul_283)
  %input.11 = Add(%layers.0.linear2.bias, %onnx::Add_139)
  %input.15 = Add(%onnx::MatMul_133, %input.11)
  %onnx::Sub_142 = ReduceMean[axes = [-1]](%input.15)
  %onnx::Pow_143 = Sub(%input.15, %onnx::Sub_142)
  %onnx::Pow_144 = Constant[value = <Scalar Tensor []>]()
  %onnx::ReduceMean_145 = Pow(%onnx::Pow_143, %onnx::Pow_144)
  %onnx::Add_146 = ReduceMean[axes = [-1]](%onnx::ReduceMean_145)
  %onnx::Add_147 = Constant[value = <Scalar Tensor []>]()
  %onnx::Sqrt_148 = Add(%onnx::Add_146, %onnx::Add_147)
  %onnx::Div_149 = Sqrt(%onnx::Sqrt_148)
  %onnx::Mul_150 = Div(%onnx::Pow_143, %onnx::Div_149)
  %onnx::Add_151 = Mul(%onnx::Mul_150, %layers.0.norm2.weight)
  %query = Add(%onnx::Add_151, %layers.0.norm2.bias)
  %onnx::Gather_153 = Shape(%query)
  %onnx::Gather_154 = Constant[value = <Scalar Tensor []>]()
  %onnx::Unsqueeze_155 = Gather[axis = 0](%onnx::Gather_153, %onnx::Gather_154)
  %onnx::Gather_156 = Shape(%query)
  %onnx::Gather_157 = Constant[value = <Scalar Tensor []>]()
  %onnx::Mul_158 = Gather[axis = 0](%onnx::Gather_156, %onnx::Gather_157)
  %onnx::Gather_159 = Shape(%query)
  %onnx::Gather_160 = Constant[value = <Scalar Tensor []>]()
  %onnx::Div_161 = Gather[axis = 0](%onnx::Gather_159, %onnx::Gather_160)
  %onnx::Div_162 = Constant[value = <Scalar Tensor []>]()
  %onnx::Cast_163 = Div(%onnx::Div_161, %onnx::Div_162)
  %onnx::Cast_164 = Cast[to = 7](%onnx::Cast_163)
  %onnx::Unsqueeze_165 = Cast[to = 7](%onnx::Cast_164)
  %onnx::Add_167 = MatMul(%query, %onnx::MatMul_284)
  %onnx::Shape_168 = Add(%layers.1.self_attn.in_proj_bias, %onnx::Add_167)
  %onnx::Gather_169 = Shape(%onnx::Shape_168)
  %onnx::Gather_170 = Constant[value = <Tensor>]()
  %onnx::Add_171 = Gather[axis = 0](%onnx::Gather_169, %onnx::Gather_170)
  %onnx::Slice_172 = Constant[value = <Tensor>]()
  %onnx::Add_173 = Constant[value = <Tensor>]()
  %onnx::Div_174 = Add(%onnx::Add_171, %onnx::Add_173)
  %onnx::Div_175 = Constant[value = <Tensor>]()
  %onnx::Mul_176 = Div(%onnx::Div_174, %onnx::Div_175)
  %onnx::Mul_177 = Constant[value = <Tensor>]()
  %onnx::Slice_178 = Mul(%onnx::Mul_176, %onnx::Mul_177)
  %onnx::Reshape_179 = Slice(%onnx::Shape_168, %onnx::Slice_172, %onnx::Slice_178, %onnx::Gather_170)
  %onnx::Mul_180 = Constant[value = <Tensor>]()
  %onnx::Slice_181 = Mul(%onnx::Mul_176, %onnx::Mul_180)
  %onnx::Shape_182 = Slice(%onnx::Shape_168, %onnx::Slice_178, %onnx::Slice_181, %onnx::Gather_170)
  %onnx::Mul_183 = Constant[value = <Tensor>]()
  %onnx::Slice_184 = Mul(%onnx::Mul_176, %onnx::Mul_183)
  %onnx::Shape_185 = Slice(%onnx::Shape_168, %onnx::Slice_181, %onnx::Slice_184, %onnx::Gather_170)
  %onnx::Mul_186 = Constant[value = <Scalar Tensor []>]()
  %onnx::Unsqueeze_187 = Mul(%onnx::Mul_158, %onnx::Mul_186)
  %onnx::Unsqueeze_188 = Constant[value = <Tensor>]()
  %onnx::Concat_189 = Unsqueeze(%onnx::Unsqueeze_155, %onnx::Unsqueeze_188)
  %onnx::Unsqueeze_190 = Constant[value = <Tensor>]()
  %onnx::Concat_191 = Unsqueeze(%onnx::Unsqueeze_187, %onnx::Unsqueeze_190)
  %onnx::Unsqueeze_192 = Constant[value = <Tensor>]()
  %onnx::Concat_193 = Unsqueeze(%onnx::Unsqueeze_165, %onnx::Unsqueeze_192)
  %onnx::Reshape_194 = Concat[axis = 0](%onnx::Concat_189, %onnx::Concat_191, %onnx::Concat_193)
  %onnx::Transpose_195 = Reshape(%onnx::Reshape_179, %onnx::Reshape_194)
  %q.3 = Transpose[perm = [1, 0, 2]](%onnx::Transpose_195)
  %onnx::Gather_197 = Shape(%onnx::Shape_182)
  %onnx::Gather_198 = Constant[value = <Scalar Tensor []>]()
  %onnx::Unsqueeze_199 = Gather[axis = 0](%onnx::Gather_197, %onnx::Gather_198)
  %onnx::Mul_200 = Constant[value = <Scalar Tensor []>]()
  %onnx::Unsqueeze_201 = Mul(%onnx::Mul_158, %onnx::Mul_200)
  %onnx::Unsqueeze_202 = Constant[value = <Tensor>]()
  %onnx::Concat_203 = Unsqueeze(%onnx::Unsqueeze_199, %onnx::Unsqueeze_202)
  %onnx::Unsqueeze_204 = Constant[value = <Tensor>]()
  %onnx::Concat_205 = Unsqueeze(%onnx::Unsqueeze_201, %onnx::Unsqueeze_204)
  %onnx::Unsqueeze_206 = Constant[value = <Tensor>]()
  %onnx::Concat_207 = Unsqueeze(%onnx::Unsqueeze_165, %onnx::Unsqueeze_206)
  %onnx::Reshape_208 = Concat[axis = 0](%onnx::Concat_203, %onnx::Concat_205, %onnx::Concat_207)
  %onnx::Transpose_209 = Reshape(%onnx::Shape_182, %onnx::Reshape_208)
  %onnx::Gather_210 = Shape(%onnx::Shape_185)
  %onnx::Gather_211 = Constant[value = <Scalar Tensor []>]()
  %onnx::Unsqueeze_212 = Gather[axis = 0](%onnx::Gather_210, %onnx::Gather_211)
  %onnx::Mul_213 = Constant[value = <Scalar Tensor []>]()
  %onnx::Unsqueeze_214 = Mul(%onnx::Mul_158, %onnx::Mul_213)
  %onnx::Unsqueeze_215 = Constant[value = <Tensor>]()
  %onnx::Concat_216 = Unsqueeze(%onnx::Unsqueeze_212, %onnx::Unsqueeze_215)
  %onnx::Unsqueeze_217 = Constant[value = <Tensor>]()
  %onnx::Concat_218 = Unsqueeze(%onnx::Unsqueeze_214, %onnx::Unsqueeze_217)
  %onnx::Unsqueeze_219 = Constant[value = <Tensor>]()
  %onnx::Concat_220 = Unsqueeze(%onnx::Unsqueeze_165, %onnx::Unsqueeze_219)
  %onnx::Reshape_221 = Concat[axis = 0](%onnx::Concat_216, %onnx::Concat_218, %onnx::Concat_220)
  %onnx::Transpose_222 = Reshape(%onnx::Shape_185, %onnx::Reshape_221)
  %v.3 = Transpose[perm = [1, 0, 2]](%onnx::Transpose_222)
  %onnx::Div_224 = Constant[value = <Scalar Tensor []>]()
  %onnx::MatMul_225 = Div(%q.3, %onnx::Div_224)
  %onnx::MatMul_226 = Transpose[perm = [1, 2, 0]](%onnx::Transpose_209)
  %attn.3 = MatMul(%onnx::MatMul_225, %onnx::MatMul_226)
  %onnx::MatMul_228 = Softmax[axis = -1](%attn.3)
  %onnx::Transpose_229 = MatMul(%onnx::MatMul_228, %v.3)
  %onnx::Reshape_230 = Transpose[perm = [1, 0, 2]](%onnx::Transpose_229)
  %onnx::Unsqueeze_231 = Mul(%onnx::Unsqueeze_155, %onnx::Mul_158)
  %onnx::Unsqueeze_232 = Constant[value = <Tensor>]()
  %onnx::Concat_233 = Unsqueeze(%onnx::Unsqueeze_231, %onnx::Unsqueeze_232)
  %onnx::Unsqueeze_234 = Constant[value = <Tensor>]()
  %onnx::Concat_235 = Unsqueeze(%onnx::Div_161, %onnx::Unsqueeze_234)
  %onnx::Reshape_236 = Concat[axis = 0](%onnx::Concat_233, %onnx::Concat_235)
  %onnx::Gemm_237 = Reshape(%onnx::Reshape_230, %onnx::Reshape_236)
  %onnx::Shape_238 = Gemm[alpha = 1, beta = 1, transB = 1](%onnx::Gemm_237, %layers.1.self_attn.out_proj.weight, %layers.1.self_attn.out_proj.bias)
  %onnx::Gather_239 = Shape(%onnx::Shape_238)
  %onnx::Gather_240 = Constant[value = <Scalar Tensor []>]()
  %onnx::Unsqueeze_241 = Gather[axis = 0](%onnx::Gather_239, %onnx::Gather_240)
  %onnx::Unsqueeze_242 = Constant[value = <Tensor>]()
  %onnx::Concat_243 = Unsqueeze(%onnx::Unsqueeze_155, %onnx::Unsqueeze_242)
  %onnx::Unsqueeze_244 = Constant[value = <Tensor>]()
  %onnx::Concat_245 = Unsqueeze(%onnx::Mul_158, %onnx::Unsqueeze_244)
  %onnx::Unsqueeze_246 = Constant[value = <Tensor>]()
  %onnx::Concat_247 = Unsqueeze(%onnx::Unsqueeze_241, %onnx::Unsqueeze_246)
  %onnx::Reshape_248 = Concat[axis = 0](%onnx::Concat_243, %onnx::Concat_245, %onnx::Concat_247)
  %input.19 = Reshape(%onnx::Shape_238, %onnx::Reshape_248)
  %input.23 = Add(%query, %input.19)
  %onnx::Sub_251 = ReduceMean[axes = [-1]](%input.23)
  %onnx::Pow_252 = Sub(%input.23, %onnx::Sub_251)
  %onnx::Pow_253 = Constant[value = <Scalar Tensor []>]()
  %onnx::ReduceMean_254 = Pow(%onnx::Pow_252, %onnx::Pow_253)
  %onnx::Add_255 = ReduceMean[axes = [-1]](%onnx::ReduceMean_254)
  %onnx::Add_256 = Constant[value = <Scalar Tensor []>]()
  %onnx::Sqrt_257 = Add(%onnx::Add_255, %onnx::Add_256)
  %onnx::Div_258 = Sqrt(%onnx::Sqrt_257)
  %onnx::Mul_259 = Div(%onnx::Pow_252, %onnx::Div_258)
  %onnx::Add_260 = Mul(%onnx::Mul_259, %layers.1.norm1.weight)
  %onnx::MatMul_261 = Add(%onnx::Add_260, %layers.1.norm1.bias)
  %onnx::Add_263 = MatMul(%onnx::MatMul_261, %onnx::MatMul_285)
  %onnx::Relu_264 = Add(%layers.1.linear1.bias, %onnx::Add_263)
  %input.27 = Relu(%onnx::Relu_264)
  %onnx::Add_267 = MatMul(%input.27, %onnx::MatMul_286)
  %input.31 = Add(%layers.1.linear2.bias, %onnx::Add_267)
  %input.35 = Add(%onnx::MatMul_261, %input.31)
  %onnx::Sub_270 = ReduceMean[axes = [-1]](%input.35)
  %onnx::Pow_271 = Sub(%input.35, %onnx::Sub_270)
  %onnx::Pow_272 = Constant[value = <Scalar Tensor []>]()
  %onnx::ReduceMean_273 = Pow(%onnx::Pow_271, %onnx::Pow_272)
  %onnx::Add_274 = ReduceMean[axes = [-1]](%onnx::ReduceMean_273)
  %onnx::Add_275 = Constant[value = <Scalar Tensor []>]()
  %onnx::Sqrt_276 = Add(%onnx::Add_274, %onnx::Add_275)
  %onnx::Div_277 = Sqrt(%onnx::Sqrt_276)
  %onnx::Mul_278 = Div(%onnx::Pow_271, %onnx::Div_277)
  %onnx::Add_279 = Mul(%onnx::Mul_278, %layers.1.norm2.weight)
  %Output = Add(%onnx::Add_279, %layers.1.norm2.bias)
  return %Output
}