test_dag.mlir

module {
  func.func @main_graph(%arg0: !torch.vtensor<[?,?],si64>, %arg1: !torch.vtensor<[?,?],si64>, %arg2: !torch.vtensor<[?,?],si64>, %arg3: !torch.vtensor<[?,?,?],si64>, %arg4: !torch.vtensor<[4],si64>, %arg5: !torch.vtensor<[2],si64>) -> !torch.vtensor<[?,?,128,384],i1>   attributes {torch.onnx_meta.ir_version = 7 : si64, torch.onnx_meta.opset_version = 21 : si64, torch.onnx_meta.producer_name = "pytorch", torch.onnx_meta.producer_version = "2.4.0"} {
    %1 = torch.operator "onnx.Constant"() {torch.onnx.value = dense<0> : tensor<si64>} : () -> !torch.vtensor<[],si64>
    %2 = torch.operator "onnx.Pad"(%arg1, %arg4, %1) {torch.onnx.mode = "constant"} : (!torch.vtensor<[?,?],si64>, !torch.vtensor<[4],si64>, !torch.vtensor<[],si64>) -> !torch.vtensor<[?,?],si64>
    %3 = torch.operator "onnx.Constant"() {torch.onnx.value = dense<1> : tensor<si64>} : () -> !torch.vtensor<[],si64>
    %4 = torch.operator "onnx.Gather"(%arg5, %3) {torch.onnx.axis = 0 : si64} : (!torch.vtensor<[2],si64>, !torch.vtensor<[],si64>) -> !torch.vtensor<[],si64>
    %5 = torch.operator "onnx.Constant"() {torch.onnx.value = dense<0> : tensor<si64>} : () -> !torch.vtensor<[],si64>
    %6 = torch.operator "onnx.Div"(%4, %5) : (!torch.vtensor<[],si64>, !torch.vtensor<[],si64>) -> !torch.vtensor<[],si64>
    %7 = torch.operator "onnx.Constant"() {torch.onnx.value = dense<0> : tensor<1xsi64>} : () -> !torch.vtensor<[1],si64>
    %8 = torch.operator "onnx.Constant"() {torch.onnx.value = dense<0> : tensor<1xsi64>} : () -> !torch.vtensor<[1],si64>
    %9 = torch.operator "onnx.Unsqueeze"(%6, %8) : (!torch.vtensor<[],si64>, !torch.vtensor<[1],si64>) -> !torch.vtensor<[1],si64>
    %10 = torch.operator "onnx.Constant"() {torch.onnx.value = dense<8> : tensor<1xsi64>} : () -> !torch.vtensor<[1],si64>
    %11 = torch.operator "onnx.Concat"(%7, %9, %10) {torch.onnx.axis = 0 : si64} : (!torch.vtensor<[1],si64>, !torch.vtensor<[1],si64>, !torch.vtensor<[1],si64>) -> !torch.vtensor<[3],si64>
    %12 = torch.operator "onnx.Reshape"(%2, %11) {torch.onnx.allowzero = 0 : si64} : (!torch.vtensor<[?,?],si64>, !torch.vtensor<[3],si64>) -> !torch.vtensor<[?,?,?],si64>
    %13 = torch.operator "onnx.Constant"() {torch.onnx.value = dense<1> : tensor<1xsi64>} : () -> !torch.vtensor<[1],si64>
    %14 = torch.operator "onnx.Constant"() {torch.onnx.value = dense<2> : tensor<1xsi64>} : () -> !torch.vtensor<[1],si64>
    %15 = torch.operator "onnx.Constant"() {torch.onnx.value = dense<0> : tensor<1xsi64>} : () -> !torch.vtensor<[1],si64>
    %16 = torch.operator "onnx.Constant"() {torch.onnx.value = dense<1> : tensor<1xsi64>} : () -> !torch.vtensor<[1],si64>
    %17 = torch.operator "onnx.Constant"() {torch.onnx.value = dense<-1> : tensor<1xsi64>} : () -> !torch.vtensor<[1],si64>
    %18 = torch.operator "onnx.Unsqueeze"(%12, %17) : (!torch.vtensor<[?,?,?],si64>, !torch.vtensor<[1],si64>) -> !torch.vtensor<[?,?,?,1],si64>
    %19 = torch.operator "onnx.Constant"() {torch.onnx.value = dense<2> : tensor<1xsi64>} : () -> !torch.vtensor<[1],si64>
    %20 = torch.operator "onnx.Unsqueeze"(%arg3, %19) : (!torch.vtensor<[?,?,?],si64>, !torch.vtensor<[1],si64>) -> !torch.vtensor<[?,?,1,?],si64>
    %21 = torch.operator "onnx.Cast"(%18) {torch.onnx.to = 9 : si64} : (!torch.vtensor<[?,?,?,1],si64>) -> !torch.vtensor<[?,?,?,1],i1>
    %22 = torch.operator "onnx.Cast"(%20) {torch.onnx.to = 9 : si64} : (!torch.vtensor<[?,?,1,?],si64>) -> !torch.vtensor<[?,?,1,?],i1>
    %23 = torch.operator "onnx.And"(%21, %22) : (!torch.vtensor<[?,?,?,1],i1>, !torch.vtensor<[?,?,1,?],i1>) -> !torch.vtensor<[?,?,?,?],i1>
    %24 = torch.operator "onnx.Cast"(%23) {torch.onnx.to = 9 : si64} : (!torch.vtensor<[?,?,?,?],i1>) -> !torch.vtensor<[?,?,?,?],i1>
    %25 = torch.operator "onnx.Constant"() {torch.onnx.value = dense<0> : tensor<1x1x128x384xi1>} : () -> !torch.vtensor<[1,1,128,384],i1>
    %26 = torch.operator "onnx.And"(%24, %25) : (!torch.vtensor<[?,?,?,?],i1>, !torch.vtensor<[1,1,128,384],i1>) -> !torch.vtensor<[?,?,128,384],i1>
    return %26 : !torch.vtensor<[?,?,128,384],i1>
  }
}