| // An example LSTM exported from a python reference model with dummy weights. |
| |
| // RUN: iree-run-mlir %s --target_backends=interpreter-bytecode --input_values="1x5xf32=[0 1 0 3 4]\n1x5x2x2xf32=[1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20]" --noexport_all --noprint_mlir | FileCheck %s --implicit-check-not="[" --implicit-check-not="]" --dump-input=fail |
| |
| // Exported via the XLA HLO Importer |
| // The resulting MLIR was modified by hand by changing all large constants to be |
| // splats of 0.42, removing the leading "module" wrapper, removing "name" |
| // attributes, removing extraneous 0s from float constants, and cleaning up |
| // extra whitespace. |
| |
| func @Min_reduction.47(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> { |
| %0 = xla_hlo.min %arg0, %arg1 : tensor<f32> |
| return %0 : tensor<f32> |
| } |
| func @Max_reduction.51(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> { |
| %0 = xla_hlo.max %arg0, %arg1 : tensor<i32> |
| return %0 : tensor<i32> |
| } |
| func @Max_1_reduction.55(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> { |
| %0 = xla_hlo.max %arg0, %arg1 : tensor<i32> |
| return %0 : tensor<i32> |
| } |
| func @ForwardLoopBody_o0Jnom3Cdxo__.59(%arg0: tuple<tensor<i64>, tensor<i64>, tensor<40xf32>, tensor<i64>, tensor<74x40xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>, tensor<5x1x64xf32>, tensor<5x1x1xf32>, tensor<5x1x1xf32>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>>) -> tuple<tensor<i64>, tensor<i64>, tensor<40xf32>, tensor<i64>, tensor<74x40xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>, tensor<5x1x64xf32>, tensor<5x1x1xf32>, tensor<5x1x1xf32>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>> { |
| %0 = "xla_hlo.get_tuple_element"(%arg0) {index = 0 : i32} : (tuple<tensor<i64>, tensor<i64>, tensor<40xf32>, tensor<i64>, tensor<74x40xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>, tensor<5x1x64xf32>, tensor<5x1x1xf32>, tensor<5x1x1xf32>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>>) -> tensor<i64> |
| %cst = constant dense<1> : tensor<i64> |
| %1 = xla_hlo.add %0, %cst : tensor<i64> |
| %2 = "xla_hlo.get_tuple_element"(%arg0) {index = 1 : i32} : (tuple<tensor<i64>, tensor<i64>, tensor<40xf32>, tensor<i64>, tensor<74x40xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>, tensor<5x1x64xf32>, tensor<5x1x1xf32>, tensor<5x1x1xf32>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>>) -> tensor<i64> |
| %3 = "xla_hlo.get_tuple_element"(%arg0) {index = 2 : i32} : (tuple<tensor<i64>, tensor<i64>, tensor<40xf32>, tensor<i64>, tensor<74x40xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>, tensor<5x1x64xf32>, tensor<5x1x1xf32>, tensor<5x1x1xf32>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>>) -> tensor<40xf32> |
| %4 = "xla_hlo.get_tuple_element"(%arg0) {index = 3 : i32} : (tuple<tensor<i64>, tensor<i64>, tensor<40xf32>, tensor<i64>, tensor<74x40xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>, tensor<5x1x64xf32>, tensor<5x1x1xf32>, tensor<5x1x1xf32>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>>) -> tensor<i64> |
| %5 = "xla_hlo.get_tuple_element"(%arg0) {index = 4 : i32} : (tuple<tensor<i64>, tensor<i64>, tensor<40xf32>, tensor<i64>, tensor<74x40xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>, tensor<5x1x64xf32>, tensor<5x1x1xf32>, tensor<5x1x1xf32>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>>) -> tensor<74x40xf32> |
| %6 = "xla_hlo.get_tuple_element"(%arg0) {index = 5 : i32} : (tuple<tensor<i64>, tensor<i64>, tensor<40xf32>, tensor<i64>, tensor<74x40xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>, tensor<5x1x64xf32>, tensor<5x1x1xf32>, tensor<5x1x1xf32>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>>) -> tensor<i64> |
| %7 = "xla_hlo.get_tuple_element"(%arg0) {index = 9 : i32} : (tuple<tensor<i64>, tensor<i64>, tensor<40xf32>, tensor<i64>, tensor<74x40xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>, tensor<5x1x64xf32>, tensor<5x1x1xf32>, tensor<5x1x1xf32>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>>) -> tensor<5x1x1xf32> |
| %8 = "xla_hlo.gather"(%7, %0) {collapsed_slice_dims = dense<0> : tensor<1xi64>, index_vector_dim = 0 : i64, offset_dims = dense<[0, 1]> : tensor<2xi64>, slice_sizes = dense<1> : tensor<3xi64>, start_index_map = dense<0> : tensor<1xi64>} : (tensor<5x1x1xf32>, tensor<i64>) -> tensor<1x1xf32> |
| %9 = "xla_hlo.reshape"(%8) : (tensor<1x1xf32>) -> tensor<1xf32> |
| %10 = "xla_hlo.broadcast_in_dim"(%9) {broadcast_dimensions = dense<0> : tensor<1xi64>} : (tensor<1xf32>) -> tensor<1x10xf32> |
| %cst_0 = constant dense<1.0> : tensor<f32> |
| %11 = "xla_hlo.broadcast_in_dim"(%cst_0) : (tensor<f32>) -> tensor<1x10xf32> |
| %12 = xla_hlo.mul %10, %11 : tensor<1x10xf32> |
| %cst_1 = constant dense<0.0> : tensor<f32> |
| %13 = "xla_hlo.broadcast_in_dim"(%cst_1) : (tensor<f32>) -> tensor<1x10xf32> |
| %14 = "xla_hlo.compare"(%12, %13) {comparison_direction = "GT"} : (tensor<1x10xf32>, tensor<1x10xf32>) -> tensor<1x10xi1> |
| %15 = "xla_hlo.get_tuple_element"(%arg0) {index = 6 : i32} : (tuple<tensor<i64>, tensor<i64>, tensor<40xf32>, tensor<i64>, tensor<74x40xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>, tensor<5x1x64xf32>, tensor<5x1x1xf32>, tensor<5x1x1xf32>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>>) -> tensor<1x10xf32> |
| %cst_2 = constant dense<5.0e-01> : tensor<f32> |
| %16 = "xla_hlo.broadcast_in_dim"(%cst_2) : (tensor<f32>) -> tensor<1x10xf32> |
| %17 = "xla_hlo.broadcast_in_dim"(%cst_2) : (tensor<f32>) -> tensor<1x10xf32> |
| %18 = "xla_hlo.broadcast_in_dim"(%cst_2) : (tensor<f32>) -> tensor<1x10xf32> |
| %19 = "xla_hlo.get_tuple_element"(%arg0) {index = 8 : i32} : (tuple<tensor<i64>, tensor<i64>, tensor<40xf32>, tensor<i64>, tensor<74x40xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>, tensor<5x1x64xf32>, tensor<5x1x1xf32>, tensor<5x1x1xf32>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>>) -> tensor<5x1x64xf32> |
| %20 = "xla_hlo.gather"(%19, %0) {collapsed_slice_dims = dense<0> : tensor<1xi64>, index_vector_dim = 0 : i64, offset_dims = dense<[0, 1]> : tensor<2xi64>, slice_sizes = dense<[1, 1, 64]> : tensor<3xi64>, start_index_map = dense<0> : tensor<1xi64>} : (tensor<5x1x64xf32>, tensor<i64>) -> tensor<1x64xf32> |
| %21 = "xla_hlo.get_tuple_element"(%arg0) {index = 7 : i32} : (tuple<tensor<i64>, tensor<i64>, tensor<40xf32>, tensor<i64>, tensor<74x40xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>, tensor<5x1x64xf32>, tensor<5x1x1xf32>, tensor<5x1x1xf32>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>>) -> tensor<1x10xf32> |
| %22 = "xla_hlo.concatenate"(%20, %21) {dimension = 1 : i64} : (tensor<1x64xf32>, tensor<1x10xf32>) -> tensor<1x74xf32> |
| %23 = "xla_hlo.dot"(%22, %5) {precision_config = ["DEFAULT", "DEFAULT"]} : (tensor<1x74xf32>, tensor<74x40xf32>) -> tensor<1x40xf32> |
| %24 = "xla_hlo.transpose"(%23) {permutation = dense<[0, 1]> : tensor<2xi64>} : (tensor<1x40xf32>) -> tensor<1x40xf32> |
| %25 = "xla_hlo.reshape"(%3) : (tensor<40xf32>) -> tensor<1x40xf32> |
| %26 = xla_hlo.add %24, %25 : tensor<1x40xf32> |
| %27 = "xla_hlo.slice"(%26) {limit_indices = dense<[1, 30]> : tensor<2xi64>, start_indices = dense<[0, 20]> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<1x40xf32>) -> tensor<1x10xf32> |
| %28 = xla_hlo.mul %18, %27 : tensor<1x10xf32> |
| %29 = "xla_hlo.tanh"(%28) : (tensor<1x10xf32>) -> tensor<1x10xf32> |
| %30 = xla_hlo.mul %17, %29 : tensor<1x10xf32> |
| %31 = xla_hlo.add %16, %30 : tensor<1x10xf32> |
| %32 = xla_hlo.mul %31, %15 : tensor<1x10xf32> |
| %cst_3 = constant dense<5.0e-01> : tensor<f32> |
| %33 = "xla_hlo.broadcast_in_dim"(%cst_3) : (tensor<f32>) -> tensor<1x10xf32> |
| %34 = "xla_hlo.broadcast_in_dim"(%cst_3) : (tensor<f32>) -> tensor<1x10xf32> |
| %35 = "xla_hlo.broadcast_in_dim"(%cst_3) : (tensor<f32>) -> tensor<1x10xf32> |
| %36 = "xla_hlo.slice"(%26) {limit_indices = dense<[1, 20]> : tensor<2xi64>, start_indices = dense<[0, 10]> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<1x40xf32>) -> tensor<1x10xf32> |
| %37 = xla_hlo.mul %35, %36 : tensor<1x10xf32> |
| %38 = "xla_hlo.tanh"(%37) : (tensor<1x10xf32>) -> tensor<1x10xf32> |
| %39 = xla_hlo.mul %34, %38 : tensor<1x10xf32> |
| %40 = xla_hlo.add %33, %39 : tensor<1x10xf32> |
| %41 = "xla_hlo.slice"(%26) {limit_indices = dense<[1, 10]> : tensor<2xi64>, start_indices = dense<0> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<1x40xf32>) -> tensor<1x10xf32> |
| %42 = "xla_hlo.tanh"(%41) : (tensor<1x10xf32>) -> tensor<1x10xf32> |
| %43 = xla_hlo.mul %40, %42 : tensor<1x10xf32> |
| %44 = xla_hlo.add %32, %43 : tensor<1x10xf32> |
| %cst_4 = constant dense<1.0e+01> : tensor<f32> |
| %45 = "xla_hlo.broadcast_in_dim"(%cst_4) : (tensor<f32>) -> tensor<1x10xf32> |
| %46 = xla_hlo.min %44, %45 : tensor<1x10xf32> |
| %cst_5 = constant dense<-1.0e+01> : tensor<f32> |
| %47 = "xla_hlo.broadcast_in_dim"(%cst_5) : (tensor<f32>) -> tensor<1x10xf32> |
| %48 = xla_hlo.max %46, %47 : tensor<1x10xf32> |
| %49 = "xla_hlo.select"(%14, %15, %48) : (tensor<1x10xi1>, tensor<1x10xf32>, tensor<1x10xf32>) -> tensor<1x10xf32> |
| %50 = "xla_hlo.reshape"(%8) : (tensor<1x1xf32>) -> tensor<1xf32> |
| %51 = "xla_hlo.broadcast_in_dim"(%50) {broadcast_dimensions = dense<0> : tensor<1xi64>} : (tensor<1xf32>) -> tensor<1x10xf32> |
| %cst_6 = constant dense<1.0> : tensor<f32> |
| %52 = "xla_hlo.broadcast_in_dim"(%cst_6) : (tensor<f32>) -> tensor<1x10xf32> |
| %53 = xla_hlo.mul %51, %52 : tensor<1x10xf32> |
| %cst_7 = constant dense<0.0> : tensor<f32> |
| %54 = "xla_hlo.broadcast_in_dim"(%cst_7) : (tensor<f32>) -> tensor<1x10xf32> |
| %55 = "xla_hlo.compare"(%53, %54) {comparison_direction = "GT"} : (tensor<1x10xf32>, tensor<1x10xf32>) -> tensor<1x10xi1> |
| %cst_8 = constant dense<5.0e-01> : tensor<f32> |
| %56 = "xla_hlo.broadcast_in_dim"(%cst_8) : (tensor<f32>) -> tensor<1x10xf32> |
| %57 = "xla_hlo.broadcast_in_dim"(%cst_8) : (tensor<f32>) -> tensor<1x10xf32> |
| %58 = "xla_hlo.broadcast_in_dim"(%cst_8) : (tensor<f32>) -> tensor<1x10xf32> |
| %59 = "xla_hlo.slice"(%26) {limit_indices = dense<[1, 40]> : tensor<2xi64>, start_indices = dense<[0, 30]> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<1x40xf32>) -> tensor<1x10xf32> |
| %60 = xla_hlo.mul %58, %59 : tensor<1x10xf32> |
| %61 = "xla_hlo.tanh"(%60) : (tensor<1x10xf32>) -> tensor<1x10xf32> |
| %62 = xla_hlo.mul %57, %61 : tensor<1x10xf32> |
| %63 = xla_hlo.add %56, %62 : tensor<1x10xf32> |
| %64 = "xla_hlo.tanh"(%48) : (tensor<1x10xf32>) -> tensor<1x10xf32> |
| %65 = xla_hlo.mul %63, %64 : tensor<1x10xf32> |
| %66 = "xla_hlo.select"(%55, %21, %65) : (tensor<1x10xi1>, tensor<1x10xf32>, tensor<1x10xf32>) -> tensor<1x10xf32> |
| %67 = "xla_hlo.get_tuple_element"(%arg0) {index = 10 : i32} : (tuple<tensor<i64>, tensor<i64>, tensor<40xf32>, tensor<i64>, tensor<74x40xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>, tensor<5x1x64xf32>, tensor<5x1x1xf32>, tensor<5x1x1xf32>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>>) -> tensor<5x1x1xf32> |
| %68 = "xla_hlo.get_tuple_element"(%arg0) {index = 11 : i32} : (tuple<tensor<i64>, tensor<i64>, tensor<40xf32>, tensor<i64>, tensor<74x40xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>, tensor<5x1x64xf32>, tensor<5x1x1xf32>, tensor<5x1x1xf32>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>>) -> tensor<5xi64> |
| %69 = "xla_hlo.reshape"(%6) : (tensor<i64>) -> tensor<1xi64> |
| %70 = "xla_hlo.slice"(%69) {limit_indices = dense<1> : tensor<1xi64>, start_indices = dense<0> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<1xi64>) -> tensor<1xi64> |
| %71 = "xla_hlo.reshape"(%0) : (tensor<i64>) -> tensor<1xi64> |
| %72 = "xla_hlo.concatenate"(%71) {dimension = 0 : i64} : (tensor<1xi64>) -> tensor<1xi64> |
| %73 = "xla_hlo.convert"(%72) : (tensor<1xi64>) -> tensor<1xi32> |
| %74 = "xla_hlo.slice"(%73) {limit_indices = dense<1> : tensor<1xi64>, start_indices = dense<0> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<1xi32>) -> tensor<1xi32> |
| %75 = "xla_hlo.reshape"(%74) : (tensor<1xi32>) -> tensor<i32> |
| %76 = "xla_hlo.dynamic-update-slice"(%68, %70, %75) : (tensor<5xi64>, tensor<1xi64>, tensor<i32>) -> tensor<5xi64> |
| %77 = "xla_hlo.get_tuple_element"(%arg0) {index = 12 : i32} : (tuple<tensor<i64>, tensor<i64>, tensor<40xf32>, tensor<i64>, tensor<74x40xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>, tensor<5x1x64xf32>, tensor<5x1x1xf32>, tensor<5x1x1xf32>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>>) -> tensor<5x1x10xf32> |
| %78 = "xla_hlo.reshape"(%49) : (tensor<1x10xf32>) -> tensor<1x1x10xf32> |
| %79 = "xla_hlo.slice"(%78) {limit_indices = dense<[1, 1, 10]> : tensor<3xi64>, start_indices = dense<0> : tensor<3xi64>, strides = dense<1> : tensor<3xi64>} : (tensor<1x1x10xf32>) -> tensor<1x1x10xf32> |
| %80 = "xla_hlo.slice"(%73) {limit_indices = dense<1> : tensor<1xi64>, start_indices = dense<0> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<1xi32>) -> tensor<1xi32> |
| %81 = "xla_hlo.reshape"(%80) : (tensor<1xi32>) -> tensor<i32> |
| %cst_9 = constant dense<0> : tensor<i32> |
| %82 = "xla_hlo.dynamic-update-slice"(%77, %79, %81, %cst_9, %cst_9) : (tensor<5x1x10xf32>, tensor<1x1x10xf32>, tensor<i32>, tensor<i32>, tensor<i32>) -> tensor<5x1x10xf32> |
| %83 = "xla_hlo.get_tuple_element"(%arg0) {index = 13 : i32} : (tuple<tensor<i64>, tensor<i64>, tensor<40xf32>, tensor<i64>, tensor<74x40xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>, tensor<5x1x64xf32>, tensor<5x1x1xf32>, tensor<5x1x1xf32>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>>) -> tensor<5x1x10xf32> |
| %84 = "xla_hlo.reshape"(%66) : (tensor<1x10xf32>) -> tensor<1x1x10xf32> |
| %85 = "xla_hlo.slice"(%84) {limit_indices = dense<[1, 1, 10]> : tensor<3xi64>, start_indices = dense<0> : tensor<3xi64>, strides = dense<1> : tensor<3xi64>} : (tensor<1x1x10xf32>) -> tensor<1x1x10xf32> |
| %86 = "xla_hlo.slice"(%73) {limit_indices = dense<1> : tensor<1xi64>, start_indices = dense<0> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<1xi32>) -> tensor<1xi32> |
| %87 = "xla_hlo.reshape"(%86) : (tensor<1xi32>) -> tensor<i32> |
| %cst_10 = constant dense<0> : tensor<i32> |
| %88 = "xla_hlo.dynamic-update-slice"(%83, %85, %87, %cst_10, %cst_10) : (tensor<5x1x10xf32>, tensor<1x1x10xf32>, tensor<i32>, tensor<i32>, tensor<i32>) -> tensor<5x1x10xf32> |
| %89 = "xla_hlo.tuple"(%1, %2, %3, %4, %5, %6, %49, %66, %19, %7, %67, %76, %82, %88) : (tensor<i64>, tensor<i64>, tensor<40xf32>, tensor<i64>, tensor<74x40xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>, tensor<5x1x64xf32>, tensor<5x1x1xf32>, tensor<5x1x1xf32>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>) -> tuple<tensor<i64>, tensor<i64>, tensor<40xf32>, tensor<i64>, tensor<74x40xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>, tensor<5x1x64xf32>, tensor<5x1x1xf32>, tensor<5x1x1xf32>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>> |
| return %89 : tuple<tensor<i64>, tensor<i64>, tensor<40xf32>, tensor<i64>, tensor<74x40xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>, tensor<5x1x64xf32>, tensor<5x1x1xf32>, tensor<5x1x1xf32>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>> |
| } |
| func @ForwardLoopCond_gFAnjWGSoLs__.167(%arg0: tuple<tensor<i64>, tensor<i64>, tensor<40xf32>, tensor<i64>, tensor<74x40xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>, tensor<5x1x64xf32>, tensor<5x1x1xf32>, tensor<5x1x1xf32>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>>) -> tuple<tensor<i1>> { |
| %0 = "xla_hlo.get_tuple_element"(%arg0) {index = 0 : i32} : (tuple<tensor<i64>, tensor<i64>, tensor<40xf32>, tensor<i64>, tensor<74x40xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>, tensor<5x1x64xf32>, tensor<5x1x1xf32>, tensor<5x1x1xf32>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>>) -> tensor<i64> |
| %1 = "xla_hlo.get_tuple_element"(%arg0) {index = 1 : i32} : (tuple<tensor<i64>, tensor<i64>, tensor<40xf32>, tensor<i64>, tensor<74x40xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>, tensor<5x1x64xf32>, tensor<5x1x1xf32>, tensor<5x1x1xf32>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>>) -> tensor<i64> |
| %2 = "xla_hlo.compare"(%0, %1) {comparison_direction = "LT"} : (tensor<i64>, tensor<i64>) -> tensor<i1> |
| %3 = "xla_hlo.tuple"(%2) : (tensor<i1>) -> tuple<tensor<i1>> |
| return %3 : tuple<tensor<i1>> |
| } |
| func @cond_wrapper.185(%arg0: tuple<tensor<i64>, tensor<i64>, tensor<40xf32>, tensor<i64>, tensor<74x40xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>, tensor<5x1x64xf32>, tensor<5x1x1xf32>, tensor<5x1x1xf32>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>>) -> tensor<i1> { |
| %0 = call @ForwardLoopCond_gFAnjWGSoLs__.167(%arg0) : (tuple<tensor<i64>, tensor<i64>, tensor<40xf32>, tensor<i64>, tensor<74x40xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>, tensor<5x1x64xf32>, tensor<5x1x1xf32>, tensor<5x1x1xf32>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>>) -> tuple<tensor<i1>> |
| %1 = "xla_hlo.get_tuple_element"(%0) {index = 0 : i32} : (tuple<tensor<i1>>) -> tensor<i1> |
| return %1 : tensor<i1> |
| } |
| func @Forward_DIZIAkooG44__disable_call_shape_inference_true_.189(%arg0: tensor<1x10xf32>, %arg1: tensor<1x10xf32>, %arg2: tensor<5x1x64xf32>, %arg3: tensor<5x1x1xf32>, %arg4: tensor<5x1x1xf32>) -> tuple<tensor<i64>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>> { |
| %cst = constant dense<5> : tensor<i32> |
| %0 = "xla_hlo.convert"(%arg3) : (tensor<5x1x1xf32>) -> tensor<5x1x1xf32> |
| %cst_0 = constant dense<0x7F800000> : tensor<f32> |
| %1 = "xla_hlo.convert"(%cst_0) : (tensor<f32>) -> tensor<f32> |
| %2 = "xla_hlo.reduce"(%0, %1) ( { |
| ^bb0(%arg5: tensor<f32>, %arg6: tensor<f32>): // no predecessors |
| %42 = xla_hlo.min %arg5, %arg6 : tensor<f32> |
| "xla_hlo.return"(%42) : (tensor<f32>) -> () |
| }) {dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<5x1x1xf32>, tensor<f32>) -> tensor<5xf32> |
| %3 = "xla_hlo.convert"(%2) : (tensor<5xf32>) -> tensor<5xf32> |
| %cst_1 = constant dense<0.0> : tensor<f32> |
| %4 = "xla_hlo.broadcast_in_dim"(%cst_1) : (tensor<f32>) -> tensor<5xf32> |
| %5 = "xla_hlo.compare"(%3, %4) {comparison_direction = "EQ"} : (tensor<5xf32>, tensor<5xf32>) -> tensor<5xi1> |
| %6 = "xla_hlo.convert"(%5) : (tensor<5xi1>) -> tensor<5xi32> |
| %cst_2 = constant dense<[1, 2, 3, 4, 5]> : tensor<5xi32> |
| %7 = xla_hlo.mul %6, %cst_2 : tensor<5xi32> |
| %8 = "xla_hlo.convert"(%7) : (tensor<5xi32>) -> tensor<5xi32> |
| %cst_3 = constant dense<-2147483648> : tensor<i32> |
| %9 = "xla_hlo.convert"(%cst_3) : (tensor<i32>) -> tensor<i32> |
| %10 = "xla_hlo.reduce"(%8, %9) ( { |
| ^bb0(%arg5: tensor<i32>, %arg6: tensor<i32>): // no predecessors |
| %42 = xla_hlo.max %arg5, %arg6 : tensor<i32> |
| "xla_hlo.return"(%42) : (tensor<i32>) -> () |
| }) {dimensions = dense<0> : tensor<1xi64>} : (tensor<5xi32>, tensor<i32>) -> tensor<i32> |
| %11 = "xla_hlo.convert"(%10) : (tensor<i32>) -> tensor<i32> |
| %12 = xla_hlo.sub %cst, %11 : tensor<i32> |
| %cst_4 = constant dense<5> : tensor<i32> |
| %13 = "xla_hlo.compare"(%12, %cst_4) {comparison_direction = "EQ"} : (tensor<i32>, tensor<i32>) -> tensor<i1> |
| %cst_5 = constant dense<0> : tensor<i32> |
| %cst_6 = constant dense<5> : tensor<i32> |
| %14 = "xla_hlo.reverse"(%3) {dimensions = dense<0> : tensor<1xi64>} : (tensor<5xf32>) -> tensor<5xf32> |
| %cst_7 = constant dense<0.0> : tensor<f32> |
| %15 = "xla_hlo.broadcast_in_dim"(%cst_7) : (tensor<f32>) -> tensor<5xf32> |
| %16 = "xla_hlo.compare"(%14, %15) {comparison_direction = "EQ"} : (tensor<5xf32>, tensor<5xf32>) -> tensor<5xi1> |
| %17 = "xla_hlo.convert"(%16) : (tensor<5xi1>) -> tensor<5xi32> |
| %cst_8 = constant dense<[1, 2, 3, 4, 5]> : tensor<5xi32> |
| %18 = xla_hlo.mul %17, %cst_8 : tensor<5xi32> |
| %19 = "xla_hlo.convert"(%18) : (tensor<5xi32>) -> tensor<5xi32> |
| %cst_9 = constant dense<-2147483648> : tensor<i32> |
| %20 = "xla_hlo.convert"(%cst_9) : (tensor<i32>) -> tensor<i32> |
| %21 = "xla_hlo.reduce"(%19, %20) ( { |
| ^bb0(%arg5: tensor<i32>, %arg6: tensor<i32>): // no predecessors |
| %42 = xla_hlo.max %arg5, %arg6 : tensor<i32> |
| "xla_hlo.return"(%42) : (tensor<i32>) -> () |
| }) {dimensions = dense<0> : tensor<1xi64>} : (tensor<5xi32>, tensor<i32>) -> tensor<i32> |
| %22 = "xla_hlo.convert"(%21) : (tensor<i32>) -> tensor<i32> |
| %23 = xla_hlo.sub %cst_6, %22 : tensor<i32> |
| %24 = "xla_hlo.select"(%13, %cst_5, %23) : (tensor<i1>, tensor<i32>, tensor<i32>) -> tensor<i32> |
| %25 = "xla_hlo.convert"(%24) : (tensor<i32>) -> tensor<i64> |
| %cst_10 = constant dense<5> : tensor<i32> |
| %26 = xla_hlo.sub %cst_10, %12 : tensor<i32> |
| %27 = "xla_hlo.convert"(%26) : (tensor<i32>) -> tensor<i64> |
| %cst_11 = constant dense<0.0> : tensor<f32> |
| %28 = "xla_hlo.broadcast_in_dim"(%cst_11) : (tensor<f32>) -> tensor<40xf32> |
| %cst_12 = constant dense<0> : tensor<i64> |
| %cst_13 = constant dense<0.42> : tensor<74x40xf32> |
| %cst_14 = constant dense<0> : tensor<i64> |
| %cst_15 = constant dense<0> : tensor<i64> |
| %29 = "xla_hlo.broadcast_in_dim"(%cst_15) : (tensor<i64>) -> tensor<5xi64> |
| %cst_16 = constant dense<0.0> : tensor<f32> |
| %30 = "xla_hlo.broadcast_in_dim"(%cst_16) : (tensor<f32>) -> tensor<5x1x10xf32> |
| %cst_17 = constant dense<0.0> : tensor<f32> |
| %31 = "xla_hlo.broadcast_in_dim"(%cst_17) : (tensor<f32>) -> tensor<5x1x10xf32> |
| %32 = "xla_hlo.tuple"(%25, %27, %28, %cst_12, %cst_13, %cst_14, %arg0, %arg1, %arg2, %arg3, %arg4, %29, %30, %31) : (tensor<i64>, tensor<i64>, tensor<40xf32>, tensor<i64>, tensor<74x40xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>, tensor<5x1x64xf32>, tensor<5x1x1xf32>, tensor<5x1x1xf32>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>) -> tuple<tensor<i64>, tensor<i64>, tensor<40xf32>, tensor<i64>, tensor<74x40xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>, tensor<5x1x64xf32>, tensor<5x1x1xf32>, tensor<5x1x1xf32>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>> |
| %33 = "xla_hlo.while"(%32) {body = @ForwardLoopBody_o0Jnom3Cdxo__.59, cond = @cond_wrapper.185} : (tuple<tensor<i64>, tensor<i64>, tensor<40xf32>, tensor<i64>, tensor<74x40xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>, tensor<5x1x64xf32>, tensor<5x1x1xf32>, tensor<5x1x1xf32>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>>) -> tuple<tensor<i64>, tensor<i64>, tensor<40xf32>, tensor<i64>, tensor<74x40xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>, tensor<5x1x64xf32>, tensor<5x1x1xf32>, tensor<5x1x1xf32>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>> |
| %34 = "xla_hlo.get_tuple_element"(%33) {index = 0 : i32} : (tuple<tensor<i64>, tensor<i64>, tensor<40xf32>, tensor<i64>, tensor<74x40xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>, tensor<5x1x64xf32>, tensor<5x1x1xf32>, tensor<5x1x1xf32>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>>) -> tensor<i64> |
| %35 = "xla_hlo.get_tuple_element"(%33) {index = 11 : i32} : (tuple<tensor<i64>, tensor<i64>, tensor<40xf32>, tensor<i64>, tensor<74x40xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>, tensor<5x1x64xf32>, tensor<5x1x1xf32>, tensor<5x1x1xf32>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>>) -> tensor<5xi64> |
| %36 = "xla_hlo.get_tuple_element"(%33) {index = 12 : i32} : (tuple<tensor<i64>, tensor<i64>, tensor<40xf32>, tensor<i64>, tensor<74x40xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>, tensor<5x1x64xf32>, tensor<5x1x1xf32>, tensor<5x1x1xf32>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>>) -> tensor<5x1x10xf32> |
| %37 = "xla_hlo.get_tuple_element"(%33) {index = 13 : i32} : (tuple<tensor<i64>, tensor<i64>, tensor<40xf32>, tensor<i64>, tensor<74x40xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>, tensor<5x1x64xf32>, tensor<5x1x1xf32>, tensor<5x1x1xf32>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>>) -> tensor<5x1x10xf32> |
| %38 = "xla_hlo.get_tuple_element"(%33) {index = 5 : i32} : (tuple<tensor<i64>, tensor<i64>, tensor<40xf32>, tensor<i64>, tensor<74x40xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>, tensor<5x1x64xf32>, tensor<5x1x1xf32>, tensor<5x1x1xf32>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>>) -> tensor<i64> |
| %39 = "xla_hlo.get_tuple_element"(%33) {index = 6 : i32} : (tuple<tensor<i64>, tensor<i64>, tensor<40xf32>, tensor<i64>, tensor<74x40xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>, tensor<5x1x64xf32>, tensor<5x1x1xf32>, tensor<5x1x1xf32>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>>) -> tensor<1x10xf32> |
| %40 = "xla_hlo.get_tuple_element"(%33) {index = 7 : i32} : (tuple<tensor<i64>, tensor<i64>, tensor<40xf32>, tensor<i64>, tensor<74x40xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>, tensor<5x1x64xf32>, tensor<5x1x1xf32>, tensor<5x1x1xf32>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>>) -> tensor<1x10xf32> |
| %41 = "xla_hlo.tuple"(%34, %35, %36, %37, %38, %39, %40) : (tensor<i64>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>) -> tuple<tensor<i64>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>> |
| return %41 : tuple<tensor<i64>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>> |
| } |
| func @Min_reduction.316(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> { |
| %0 = xla_hlo.min %arg0, %arg1 : tensor<f32> |
| return %0 : tensor<f32> |
| } |
| func @Max_reduction.320(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> { |
| %0 = xla_hlo.max %arg0, %arg1 : tensor<i32> |
| return %0 : tensor<i32> |
| } |
| func @Max_1_reduction.324(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> { |
| %0 = xla_hlo.max %arg0, %arg1 : tensor<i32> |
| return %0 : tensor<i32> |
| } |
| func @ForwardLoopBody_o0Jnom3Cdxo__.328(%arg0: tuple<tensor<i64>, tensor<i64>, tensor<40xf32>, tensor<i64>, tensor<20x40xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>, tensor<5x1x10xf32>, tensor<5x1x1xf32>, tensor<5x1x1xf32>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>>) -> tuple<tensor<i64>, tensor<i64>, tensor<40xf32>, tensor<i64>, tensor<20x40xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>, tensor<5x1x10xf32>, tensor<5x1x1xf32>, tensor<5x1x1xf32>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>> { |
| %0 = "xla_hlo.get_tuple_element"(%arg0) {index = 0 : i32} : (tuple<tensor<i64>, tensor<i64>, tensor<40xf32>, tensor<i64>, tensor<20x40xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>, tensor<5x1x10xf32>, tensor<5x1x1xf32>, tensor<5x1x1xf32>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>>) -> tensor<i64> |
| %cst = constant dense<1> : tensor<i64> |
| %1 = xla_hlo.add %0, %cst : tensor<i64> |
| %2 = "xla_hlo.get_tuple_element"(%arg0) {index = 1 : i32} : (tuple<tensor<i64>, tensor<i64>, tensor<40xf32>, tensor<i64>, tensor<20x40xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>, tensor<5x1x10xf32>, tensor<5x1x1xf32>, tensor<5x1x1xf32>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>>) -> tensor<i64> |
| %3 = "xla_hlo.get_tuple_element"(%arg0) {index = 2 : i32} : (tuple<tensor<i64>, tensor<i64>, tensor<40xf32>, tensor<i64>, tensor<20x40xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>, tensor<5x1x10xf32>, tensor<5x1x1xf32>, tensor<5x1x1xf32>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>>) -> tensor<40xf32> |
| %4 = "xla_hlo.get_tuple_element"(%arg0) {index = 3 : i32} : (tuple<tensor<i64>, tensor<i64>, tensor<40xf32>, tensor<i64>, tensor<20x40xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>, tensor<5x1x10xf32>, tensor<5x1x1xf32>, tensor<5x1x1xf32>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>>) -> tensor<i64> |
| %5 = "xla_hlo.get_tuple_element"(%arg0) {index = 4 : i32} : (tuple<tensor<i64>, tensor<i64>, tensor<40xf32>, tensor<i64>, tensor<20x40xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>, tensor<5x1x10xf32>, tensor<5x1x1xf32>, tensor<5x1x1xf32>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>>) -> tensor<20x40xf32> |
| %6 = "xla_hlo.get_tuple_element"(%arg0) {index = 5 : i32} : (tuple<tensor<i64>, tensor<i64>, tensor<40xf32>, tensor<i64>, tensor<20x40xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>, tensor<5x1x10xf32>, tensor<5x1x1xf32>, tensor<5x1x1xf32>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>>) -> tensor<i64> |
| %7 = "xla_hlo.get_tuple_element"(%arg0) {index = 9 : i32} : (tuple<tensor<i64>, tensor<i64>, tensor<40xf32>, tensor<i64>, tensor<20x40xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>, tensor<5x1x10xf32>, tensor<5x1x1xf32>, tensor<5x1x1xf32>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>>) -> tensor<5x1x1xf32> |
| %8 = "xla_hlo.gather"(%7, %0) {collapsed_slice_dims = dense<0> : tensor<1xi64>, index_vector_dim = 0 : i64, offset_dims = dense<[0, 1]> : tensor<2xi64>, slice_sizes = dense<1> : tensor<3xi64>, start_index_map = dense<0> : tensor<1xi64>} : (tensor<5x1x1xf32>, tensor<i64>) -> tensor<1x1xf32> |
| %9 = "xla_hlo.reshape"(%8) : (tensor<1x1xf32>) -> tensor<1xf32> |
| %10 = "xla_hlo.broadcast_in_dim"(%9) {broadcast_dimensions = dense<0> : tensor<1xi64>} : (tensor<1xf32>) -> tensor<1x10xf32> |
| %cst_0 = constant dense<1.0> : tensor<f32> |
| %11 = "xla_hlo.broadcast_in_dim"(%cst_0) : (tensor<f32>) -> tensor<1x10xf32> |
| %12 = xla_hlo.mul %10, %11 : tensor<1x10xf32> |
| %cst_1 = constant dense<0.0> : tensor<f32> |
| %13 = "xla_hlo.broadcast_in_dim"(%cst_1) : (tensor<f32>) -> tensor<1x10xf32> |
| %14 = "xla_hlo.compare"(%12, %13) {comparison_direction = "GT"} : (tensor<1x10xf32>, tensor<1x10xf32>) -> tensor<1x10xi1> |
| %15 = "xla_hlo.get_tuple_element"(%arg0) {index = 6 : i32} : (tuple<tensor<i64>, tensor<i64>, tensor<40xf32>, tensor<i64>, tensor<20x40xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>, tensor<5x1x10xf32>, tensor<5x1x1xf32>, tensor<5x1x1xf32>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>>) -> tensor<1x10xf32> |
| %cst_2 = constant dense<5.0e-01> : tensor<f32> |
| %16 = "xla_hlo.broadcast_in_dim"(%cst_2) : (tensor<f32>) -> tensor<1x10xf32> |
| %17 = "xla_hlo.broadcast_in_dim"(%cst_2) : (tensor<f32>) -> tensor<1x10xf32> |
| %18 = "xla_hlo.broadcast_in_dim"(%cst_2) : (tensor<f32>) -> tensor<1x10xf32> |
| %19 = "xla_hlo.get_tuple_element"(%arg0) {index = 8 : i32} : (tuple<tensor<i64>, tensor<i64>, tensor<40xf32>, tensor<i64>, tensor<20x40xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>, tensor<5x1x10xf32>, tensor<5x1x1xf32>, tensor<5x1x1xf32>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>>) -> tensor<5x1x10xf32> |
| %20 = "xla_hlo.gather"(%19, %0) {collapsed_slice_dims = dense<0> : tensor<1xi64>, index_vector_dim = 0 : i64, offset_dims = dense<[0, 1]> : tensor<2xi64>, slice_sizes = dense<[1, 1, 10]> : tensor<3xi64>, start_index_map = dense<0> : tensor<1xi64>} : (tensor<5x1x10xf32>, tensor<i64>) -> tensor<1x10xf32> |
| %21 = "xla_hlo.get_tuple_element"(%arg0) {index = 7 : i32} : (tuple<tensor<i64>, tensor<i64>, tensor<40xf32>, tensor<i64>, tensor<20x40xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>, tensor<5x1x10xf32>, tensor<5x1x1xf32>, tensor<5x1x1xf32>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>>) -> tensor<1x10xf32> |
| %22 = "xla_hlo.concatenate"(%20, %21) {dimension = 1 : i64} : (tensor<1x10xf32>, tensor<1x10xf32>) -> tensor<1x20xf32> |
| %23 = "xla_hlo.dot"(%22, %5) {precision_config = ["DEFAULT", "DEFAULT"]} : (tensor<1x20xf32>, tensor<20x40xf32>) -> tensor<1x40xf32> |
| %24 = "xla_hlo.transpose"(%23) {permutation = dense<[0, 1]> : tensor<2xi64>} : (tensor<1x40xf32>) -> tensor<1x40xf32> |
| %25 = "xla_hlo.reshape"(%3) : (tensor<40xf32>) -> tensor<1x40xf32> |
| %26 = xla_hlo.add %24, %25 : tensor<1x40xf32> |
| %27 = "xla_hlo.slice"(%26) {limit_indices = dense<[1, 30]> : tensor<2xi64>, start_indices = dense<[0, 20]> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<1x40xf32>) -> tensor<1x10xf32> |
| %28 = xla_hlo.mul %18, %27 : tensor<1x10xf32> |
| %29 = "xla_hlo.tanh"(%28) : (tensor<1x10xf32>) -> tensor<1x10xf32> |
| %30 = xla_hlo.mul %17, %29 : tensor<1x10xf32> |
| %31 = xla_hlo.add %16, %30 : tensor<1x10xf32> |
| %32 = xla_hlo.mul %31, %15 : tensor<1x10xf32> |
| %cst_3 = constant dense<5.0e-01> : tensor<f32> |
| %33 = "xla_hlo.broadcast_in_dim"(%cst_3) : (tensor<f32>) -> tensor<1x10xf32> |
| %34 = "xla_hlo.broadcast_in_dim"(%cst_3) : (tensor<f32>) -> tensor<1x10xf32> |
| %35 = "xla_hlo.broadcast_in_dim"(%cst_3) : (tensor<f32>) -> tensor<1x10xf32> |
| %36 = "xla_hlo.slice"(%26) {limit_indices = dense<[1, 20]> : tensor<2xi64>, start_indices = dense<[0, 10]> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<1x40xf32>) -> tensor<1x10xf32> |
| %37 = xla_hlo.mul %35, %36 : tensor<1x10xf32> |
| %38 = "xla_hlo.tanh"(%37) : (tensor<1x10xf32>) -> tensor<1x10xf32> |
| %39 = xla_hlo.mul %34, %38 : tensor<1x10xf32> |
| %40 = xla_hlo.add %33, %39 : tensor<1x10xf32> |
| %41 = "xla_hlo.slice"(%26) {limit_indices = dense<[1, 10]> : tensor<2xi64>, start_indices = dense<0> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<1x40xf32>) -> tensor<1x10xf32> |
| %42 = "xla_hlo.tanh"(%41) : (tensor<1x10xf32>) -> tensor<1x10xf32> |
| %43 = xla_hlo.mul %40, %42 : tensor<1x10xf32> |
| %44 = xla_hlo.add %32, %43 : tensor<1x10xf32> |
| %cst_4 = constant dense<1.0e+01> : tensor<f32> |
| %45 = "xla_hlo.broadcast_in_dim"(%cst_4) : (tensor<f32>) -> tensor<1x10xf32> |
| %46 = xla_hlo.min %44, %45 : tensor<1x10xf32> |
| %cst_5 = constant dense<-1.0e+01> : tensor<f32> |
| %47 = "xla_hlo.broadcast_in_dim"(%cst_5) : (tensor<f32>) -> tensor<1x10xf32> |
| %48 = xla_hlo.max %46, %47 : tensor<1x10xf32> |
| %49 = "xla_hlo.select"(%14, %15, %48) : (tensor<1x10xi1>, tensor<1x10xf32>, tensor<1x10xf32>) -> tensor<1x10xf32> |
| %50 = "xla_hlo.reshape"(%8) : (tensor<1x1xf32>) -> tensor<1xf32> |
| %51 = "xla_hlo.broadcast_in_dim"(%50) {broadcast_dimensions = dense<0> : tensor<1xi64>} : (tensor<1xf32>) -> tensor<1x10xf32> |
| %cst_6 = constant dense<1.0> : tensor<f32> |
| %52 = "xla_hlo.broadcast_in_dim"(%cst_6) : (tensor<f32>) -> tensor<1x10xf32> |
| %53 = xla_hlo.mul %51, %52 : tensor<1x10xf32> |
| %cst_7 = constant dense<0.0> : tensor<f32> |
| %54 = "xla_hlo.broadcast_in_dim"(%cst_7) : (tensor<f32>) -> tensor<1x10xf32> |
| %55 = "xla_hlo.compare"(%53, %54) {comparison_direction = "GT"} : (tensor<1x10xf32>, tensor<1x10xf32>) -> tensor<1x10xi1> |
| %cst_8 = constant dense<5.0e-01> : tensor<f32> |
| %56 = "xla_hlo.broadcast_in_dim"(%cst_8) : (tensor<f32>) -> tensor<1x10xf32> |
| %57 = "xla_hlo.broadcast_in_dim"(%cst_8) : (tensor<f32>) -> tensor<1x10xf32> |
| %58 = "xla_hlo.broadcast_in_dim"(%cst_8) : (tensor<f32>) -> tensor<1x10xf32> |
| %59 = "xla_hlo.slice"(%26) {limit_indices = dense<[1, 40]> : tensor<2xi64>, start_indices = dense<[0, 30]> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<1x40xf32>) -> tensor<1x10xf32> |
| %60 = xla_hlo.mul %58, %59 : tensor<1x10xf32> |
| %61 = "xla_hlo.tanh"(%60) : (tensor<1x10xf32>) -> tensor<1x10xf32> |
| %62 = xla_hlo.mul %57, %61 : tensor<1x10xf32> |
| %63 = xla_hlo.add %56, %62 : tensor<1x10xf32> |
| %64 = "xla_hlo.tanh"(%48) : (tensor<1x10xf32>) -> tensor<1x10xf32> |
| %65 = xla_hlo.mul %63, %64 : tensor<1x10xf32> |
| %66 = "xla_hlo.select"(%55, %21, %65) : (tensor<1x10xi1>, tensor<1x10xf32>, tensor<1x10xf32>) -> tensor<1x10xf32> |
| %67 = "xla_hlo.get_tuple_element"(%arg0) {index = 10 : i32} : (tuple<tensor<i64>, tensor<i64>, tensor<40xf32>, tensor<i64>, tensor<20x40xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>, tensor<5x1x10xf32>, tensor<5x1x1xf32>, tensor<5x1x1xf32>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>>) -> tensor<5x1x1xf32> |
| %68 = "xla_hlo.get_tuple_element"(%arg0) {index = 11 : i32} : (tuple<tensor<i64>, tensor<i64>, tensor<40xf32>, tensor<i64>, tensor<20x40xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>, tensor<5x1x10xf32>, tensor<5x1x1xf32>, tensor<5x1x1xf32>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>>) -> tensor<5xi64> |
| %69 = "xla_hlo.reshape"(%6) : (tensor<i64>) -> tensor<1xi64> |
| %70 = "xla_hlo.slice"(%69) {limit_indices = dense<1> : tensor<1xi64>, start_indices = dense<0> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<1xi64>) -> tensor<1xi64> |
| %71 = "xla_hlo.reshape"(%0) : (tensor<i64>) -> tensor<1xi64> |
| %72 = "xla_hlo.concatenate"(%71) {dimension = 0 : i64} : (tensor<1xi64>) -> tensor<1xi64> |
| %73 = "xla_hlo.convert"(%72) : (tensor<1xi64>) -> tensor<1xi32> |
| %74 = "xla_hlo.slice"(%73) {limit_indices = dense<1> : tensor<1xi64>, start_indices = dense<0> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<1xi32>) -> tensor<1xi32> |
| %75 = "xla_hlo.reshape"(%74) : (tensor<1xi32>) -> tensor<i32> |
| %76 = "xla_hlo.dynamic-update-slice"(%68, %70, %75) : (tensor<5xi64>, tensor<1xi64>, tensor<i32>) -> tensor<5xi64> |
| %77 = "xla_hlo.get_tuple_element"(%arg0) {index = 12 : i32} : (tuple<tensor<i64>, tensor<i64>, tensor<40xf32>, tensor<i64>, tensor<20x40xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>, tensor<5x1x10xf32>, tensor<5x1x1xf32>, tensor<5x1x1xf32>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>>) -> tensor<5x1x10xf32> |
| %78 = "xla_hlo.reshape"(%49) : (tensor<1x10xf32>) -> tensor<1x1x10xf32> |
| %79 = "xla_hlo.slice"(%78) {limit_indices = dense<[1, 1, 10]> : tensor<3xi64>, start_indices = dense<0> : tensor<3xi64>, strides = dense<1> : tensor<3xi64>} : (tensor<1x1x10xf32>) -> tensor<1x1x10xf32> |
| %80 = "xla_hlo.slice"(%73) {limit_indices = dense<1> : tensor<1xi64>, start_indices = dense<0> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<1xi32>) -> tensor<1xi32> |
| %81 = "xla_hlo.reshape"(%80) : (tensor<1xi32>) -> tensor<i32> |
| %cst_9 = constant dense<0> : tensor<i32> |
| %82 = "xla_hlo.dynamic-update-slice"(%77, %79, %81, %cst_9, %cst_9) : (tensor<5x1x10xf32>, tensor<1x1x10xf32>, tensor<i32>, tensor<i32>, tensor<i32>) -> tensor<5x1x10xf32> |
| %83 = "xla_hlo.get_tuple_element"(%arg0) {index = 13 : i32} : (tuple<tensor<i64>, tensor<i64>, tensor<40xf32>, tensor<i64>, tensor<20x40xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>, tensor<5x1x10xf32>, tensor<5x1x1xf32>, tensor<5x1x1xf32>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>>) -> tensor<5x1x10xf32> |
| %84 = "xla_hlo.reshape"(%66) : (tensor<1x10xf32>) -> tensor<1x1x10xf32> |
| %85 = "xla_hlo.slice"(%84) {limit_indices = dense<[1, 1, 10]> : tensor<3xi64>, start_indices = dense<0> : tensor<3xi64>, strides = dense<1> : tensor<3xi64>} : (tensor<1x1x10xf32>) -> tensor<1x1x10xf32> |
| %86 = "xla_hlo.slice"(%73) {limit_indices = dense<1> : tensor<1xi64>, start_indices = dense<0> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<1xi32>) -> tensor<1xi32> |
| %87 = "xla_hlo.reshape"(%86) : (tensor<1xi32>) -> tensor<i32> |
| %cst_10 = constant dense<0> : tensor<i32> |
| %88 = "xla_hlo.dynamic-update-slice"(%83, %85, %87, %cst_10, %cst_10) : (tensor<5x1x10xf32>, tensor<1x1x10xf32>, tensor<i32>, tensor<i32>, tensor<i32>) -> tensor<5x1x10xf32> |
| %89 = "xla_hlo.tuple"(%1, %2, %3, %4, %5, %6, %49, %66, %19, %7, %67, %76, %82, %88) : (tensor<i64>, tensor<i64>, tensor<40xf32>, tensor<i64>, tensor<20x40xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>, tensor<5x1x10xf32>, tensor<5x1x1xf32>, tensor<5x1x1xf32>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>) -> tuple<tensor<i64>, tensor<i64>, tensor<40xf32>, tensor<i64>, tensor<20x40xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>, tensor<5x1x10xf32>, tensor<5x1x1xf32>, tensor<5x1x1xf32>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>> |
| return %89 : tuple<tensor<i64>, tensor<i64>, tensor<40xf32>, tensor<i64>, tensor<20x40xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>, tensor<5x1x10xf32>, tensor<5x1x1xf32>, tensor<5x1x1xf32>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>> |
| } |
| func @ForwardLoopCond_gFAnjWGSoLs__.436(%arg0: tuple<tensor<i64>, tensor<i64>, tensor<40xf32>, tensor<i64>, tensor<20x40xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>, tensor<5x1x10xf32>, tensor<5x1x1xf32>, tensor<5x1x1xf32>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>>) -> tuple<tensor<i1>> { |
| %0 = "xla_hlo.get_tuple_element"(%arg0) {index = 0 : i32} : (tuple<tensor<i64>, tensor<i64>, tensor<40xf32>, tensor<i64>, tensor<20x40xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>, tensor<5x1x10xf32>, tensor<5x1x1xf32>, tensor<5x1x1xf32>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>>) -> tensor<i64> |
| %1 = "xla_hlo.get_tuple_element"(%arg0) {index = 1 : i32} : (tuple<tensor<i64>, tensor<i64>, tensor<40xf32>, tensor<i64>, tensor<20x40xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>, tensor<5x1x10xf32>, tensor<5x1x1xf32>, tensor<5x1x1xf32>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>>) -> tensor<i64> |
| %2 = "xla_hlo.compare"(%0, %1) {comparison_direction = "LT"} : (tensor<i64>, tensor<i64>) -> tensor<i1> |
| %3 = "xla_hlo.tuple"(%2) : (tensor<i1>) -> tuple<tensor<i1>> |
| return %3 : tuple<tensor<i1>> |
| } |
| func @cond_wrapper.454(%arg0: tuple<tensor<i64>, tensor<i64>, tensor<40xf32>, tensor<i64>, tensor<20x40xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>, tensor<5x1x10xf32>, tensor<5x1x1xf32>, tensor<5x1x1xf32>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>>) -> tensor<i1> { |
| %0 = call @ForwardLoopCond_gFAnjWGSoLs__.436(%arg0) : (tuple<tensor<i64>, tensor<i64>, tensor<40xf32>, tensor<i64>, tensor<20x40xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>, tensor<5x1x10xf32>, tensor<5x1x1xf32>, tensor<5x1x1xf32>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>>) -> tuple<tensor<i1>> |
| %1 = "xla_hlo.get_tuple_element"(%0) {index = 0 : i32} : (tuple<tensor<i1>>) -> tensor<i1> |
| return %1 : tensor<i1> |
| } |
| func @Forward_DIZIAkooG44__disable_call_shape_inference_true_.458(%arg0: tensor<1x10xf32>, %arg1: tensor<1x10xf32>, %arg2: tensor<5x1x10xf32>, %arg3: tensor<5x1x1xf32>, %arg4: tensor<5x1x1xf32>) -> tuple<tensor<i64>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>> { |
| %cst = constant dense<5> : tensor<i32> |
| %0 = "xla_hlo.convert"(%arg3) : (tensor<5x1x1xf32>) -> tensor<5x1x1xf32> |
| %cst_0 = constant dense<0x7F800000> : tensor<f32> |
| %1 = "xla_hlo.convert"(%cst_0) : (tensor<f32>) -> tensor<f32> |
| %2 = "xla_hlo.reduce"(%0, %1) ( { |
| ^bb0(%arg5: tensor<f32>, %arg6: tensor<f32>): // no predecessors |
| %42 = xla_hlo.min %arg5, %arg6 : tensor<f32> |
| "xla_hlo.return"(%42) : (tensor<f32>) -> () |
| }) {dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<5x1x1xf32>, tensor<f32>) -> tensor<5xf32> |
| %3 = "xla_hlo.convert"(%2) : (tensor<5xf32>) -> tensor<5xf32> |
| %cst_1 = constant dense<0.0> : tensor<f32> |
| %4 = "xla_hlo.broadcast_in_dim"(%cst_1) : (tensor<f32>) -> tensor<5xf32> |
| %5 = "xla_hlo.compare"(%3, %4) {comparison_direction = "EQ"} : (tensor<5xf32>, tensor<5xf32>) -> tensor<5xi1> |
| %6 = "xla_hlo.convert"(%5) : (tensor<5xi1>) -> tensor<5xi32> |
| %cst_2 = constant dense<[1, 2, 3, 4, 5]> : tensor<5xi32> |
| %7 = xla_hlo.mul %6, %cst_2 : tensor<5xi32> |
| %8 = "xla_hlo.convert"(%7) : (tensor<5xi32>) -> tensor<5xi32> |
| %cst_3 = constant dense<-2147483648> : tensor<i32> |
| %9 = "xla_hlo.convert"(%cst_3) : (tensor<i32>) -> tensor<i32> |
| %10 = "xla_hlo.reduce"(%8, %9) ( { |
| ^bb0(%arg5: tensor<i32>, %arg6: tensor<i32>): // no predecessors |
| %42 = xla_hlo.max %arg5, %arg6 : tensor<i32> |
| "xla_hlo.return"(%42) : (tensor<i32>) -> () |
| }) {dimensions = dense<0> : tensor<1xi64>} : (tensor<5xi32>, tensor<i32>) -> tensor<i32> |
| %11 = "xla_hlo.convert"(%10) : (tensor<i32>) -> tensor<i32> |
| %12 = xla_hlo.sub %cst, %11 : tensor<i32> |
| %cst_4 = constant dense<5> : tensor<i32> |
| %13 = "xla_hlo.compare"(%12, %cst_4) {comparison_direction = "EQ"} : (tensor<i32>, tensor<i32>) -> tensor<i1> |
| %cst_5 = constant dense<0> : tensor<i32> |
| %cst_6 = constant dense<5> : tensor<i32> |
| %14 = "xla_hlo.reverse"(%3) {dimensions = dense<0> : tensor<1xi64>} : (tensor<5xf32>) -> tensor<5xf32> |
| %cst_7 = constant dense<0.0> : tensor<f32> |
| %15 = "xla_hlo.broadcast_in_dim"(%cst_7) : (tensor<f32>) -> tensor<5xf32> |
| %16 = "xla_hlo.compare"(%14, %15) {comparison_direction = "EQ"} : (tensor<5xf32>, tensor<5xf32>) -> tensor<5xi1> |
| %17 = "xla_hlo.convert"(%16) : (tensor<5xi1>) -> tensor<5xi32> |
| %cst_8 = constant dense<[1, 2, 3, 4, 5]> : tensor<5xi32> |
| %18 = xla_hlo.mul %17, %cst_8 : tensor<5xi32> |
| %19 = "xla_hlo.convert"(%18) : (tensor<5xi32>) -> tensor<5xi32> |
| %cst_9 = constant dense<-2147483648> : tensor<i32> |
| %20 = "xla_hlo.convert"(%cst_9) : (tensor<i32>) -> tensor<i32> |
| %21 = "xla_hlo.reduce"(%19, %20) ( { |
| ^bb0(%arg5: tensor<i32>, %arg6: tensor<i32>): // no predecessors |
| %42 = xla_hlo.max %arg5, %arg6 : tensor<i32> |
| "xla_hlo.return"(%42) : (tensor<i32>) -> () |
| }) {dimensions = dense<0> : tensor<1xi64>} : (tensor<5xi32>, tensor<i32>) -> tensor<i32> |
| %22 = "xla_hlo.convert"(%21) : (tensor<i32>) -> tensor<i32> |
| %23 = xla_hlo.sub %cst_6, %22 : tensor<i32> |
| %24 = "xla_hlo.select"(%13, %cst_5, %23) : (tensor<i1>, tensor<i32>, tensor<i32>) -> tensor<i32> |
| %25 = "xla_hlo.convert"(%24) : (tensor<i32>) -> tensor<i64> |
| %cst_10 = constant dense<5> : tensor<i32> |
| %26 = xla_hlo.sub %cst_10, %12 : tensor<i32> |
| %27 = "xla_hlo.convert"(%26) : (tensor<i32>) -> tensor<i64> |
| %cst_11 = constant dense<0.0> : tensor<f32> |
| %28 = "xla_hlo.broadcast_in_dim"(%cst_11) : (tensor<f32>) -> tensor<40xf32> |
| %cst_12 = constant dense<0> : tensor<i64> |
| %cst_13 = constant dense<0.42> : tensor<20x40xf32> |
| %cst_14 = constant dense<-2130569576> : tensor<i64> |
| %cst_15 = constant dense<0> : tensor<i64> |
| %29 = "xla_hlo.broadcast_in_dim"(%cst_15) : (tensor<i64>) -> tensor<5xi64> |
| %cst_16 = constant dense<0.0> : tensor<f32> |
| %30 = "xla_hlo.broadcast_in_dim"(%cst_16) : (tensor<f32>) -> tensor<5x1x10xf32> |
| %cst_17 = constant dense<0.0> : tensor<f32> |
| %31 = "xla_hlo.broadcast_in_dim"(%cst_17) : (tensor<f32>) -> tensor<5x1x10xf32> |
| %32 = "xla_hlo.tuple"(%25, %27, %28, %cst_12, %cst_13, %cst_14, %arg0, %arg1, %arg2, %arg3, %arg4, %29, %30, %31) : (tensor<i64>, tensor<i64>, tensor<40xf32>, tensor<i64>, tensor<20x40xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>, tensor<5x1x10xf32>, tensor<5x1x1xf32>, tensor<5x1x1xf32>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>) -> tuple<tensor<i64>, tensor<i64>, tensor<40xf32>, tensor<i64>, tensor<20x40xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>, tensor<5x1x10xf32>, tensor<5x1x1xf32>, tensor<5x1x1xf32>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>> |
| %33 = "xla_hlo.while"(%32) {body = @ForwardLoopBody_o0Jnom3Cdxo__.328, cond = @cond_wrapper.454} : (tuple<tensor<i64>, tensor<i64>, tensor<40xf32>, tensor<i64>, tensor<20x40xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>, tensor<5x1x10xf32>, tensor<5x1x1xf32>, tensor<5x1x1xf32>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>>) -> tuple<tensor<i64>, tensor<i64>, tensor<40xf32>, tensor<i64>, tensor<20x40xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>, tensor<5x1x10xf32>, tensor<5x1x1xf32>, tensor<5x1x1xf32>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>> |
| %34 = "xla_hlo.get_tuple_element"(%33) {index = 0 : i32} : (tuple<tensor<i64>, tensor<i64>, tensor<40xf32>, tensor<i64>, tensor<20x40xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>, tensor<5x1x10xf32>, tensor<5x1x1xf32>, tensor<5x1x1xf32>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>>) -> tensor<i64> |
| %35 = "xla_hlo.get_tuple_element"(%33) {index = 11 : i32} : (tuple<tensor<i64>, tensor<i64>, tensor<40xf32>, tensor<i64>, tensor<20x40xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>, tensor<5x1x10xf32>, tensor<5x1x1xf32>, tensor<5x1x1xf32>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>>) -> tensor<5xi64> |
| %36 = "xla_hlo.get_tuple_element"(%33) {index = 12 : i32} : (tuple<tensor<i64>, tensor<i64>, tensor<40xf32>, tensor<i64>, tensor<20x40xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>, tensor<5x1x10xf32>, tensor<5x1x1xf32>, tensor<5x1x1xf32>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>>) -> tensor<5x1x10xf32> |
| %37 = "xla_hlo.get_tuple_element"(%33) {index = 13 : i32} : (tuple<tensor<i64>, tensor<i64>, tensor<40xf32>, tensor<i64>, tensor<20x40xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>, tensor<5x1x10xf32>, tensor<5x1x1xf32>, tensor<5x1x1xf32>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>>) -> tensor<5x1x10xf32> |
| %38 = "xla_hlo.get_tuple_element"(%33) {index = 5 : i32} : (tuple<tensor<i64>, tensor<i64>, tensor<40xf32>, tensor<i64>, tensor<20x40xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>, tensor<5x1x10xf32>, tensor<5x1x1xf32>, tensor<5x1x1xf32>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>>) -> tensor<i64> |
| %39 = "xla_hlo.get_tuple_element"(%33) {index = 6 : i32} : (tuple<tensor<i64>, tensor<i64>, tensor<40xf32>, tensor<i64>, tensor<20x40xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>, tensor<5x1x10xf32>, tensor<5x1x1xf32>, tensor<5x1x1xf32>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>>) -> tensor<1x10xf32> |
| %40 = "xla_hlo.get_tuple_element"(%33) {index = 7 : i32} : (tuple<tensor<i64>, tensor<i64>, tensor<40xf32>, tensor<i64>, tensor<20x40xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>, tensor<5x1x10xf32>, tensor<5x1x1xf32>, tensor<5x1x1xf32>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>>) -> tensor<1x10xf32> |
| %41 = "xla_hlo.tuple"(%34, %35, %36, %37, %38, %39, %40) : (tensor<i64>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>) -> tuple<tensor<i64>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>> |
| return %41 : tuple<tensor<i64>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>> |
| } |
| func @Min_reduction.585(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> { |
| %0 = xla_hlo.min %arg0, %arg1 : tensor<f32> |
| return %0 : tensor<f32> |
| } |
| func @Max_reduction.589(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> { |
| %0 = xla_hlo.max %arg0, %arg1 : tensor<i32> |
| return %0 : tensor<i32> |
| } |
| func @Max_1_reduction.593(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> { |
| %0 = xla_hlo.max %arg0, %arg1 : tensor<i32> |
| return %0 : tensor<i32> |
| } |
| func @ForwardLoopBody_o0Jnom3Cdxo__.597(%arg0: tuple<tensor<i64>, tensor<i64>, tensor<40xf32>, tensor<i64>, tensor<20x40xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>, tensor<5x1x10xf32>, tensor<5x1x1xf32>, tensor<5x1x1xf32>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>>) -> tuple<tensor<i64>, tensor<i64>, tensor<40xf32>, tensor<i64>, tensor<20x40xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>, tensor<5x1x10xf32>, tensor<5x1x1xf32>, tensor<5x1x1xf32>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>> { |
| %0 = "xla_hlo.get_tuple_element"(%arg0) {index = 0 : i32} : (tuple<tensor<i64>, tensor<i64>, tensor<40xf32>, tensor<i64>, tensor<20x40xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>, tensor<5x1x10xf32>, tensor<5x1x1xf32>, tensor<5x1x1xf32>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>>) -> tensor<i64> |
| %cst = constant dense<1> : tensor<i64> |
| %1 = xla_hlo.add %0, %cst : tensor<i64> |
| %2 = "xla_hlo.get_tuple_element"(%arg0) {index = 1 : i32} : (tuple<tensor<i64>, tensor<i64>, tensor<40xf32>, tensor<i64>, tensor<20x40xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>, tensor<5x1x10xf32>, tensor<5x1x1xf32>, tensor<5x1x1xf32>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>>) -> tensor<i64> |
| %3 = "xla_hlo.get_tuple_element"(%arg0) {index = 2 : i32} : (tuple<tensor<i64>, tensor<i64>, tensor<40xf32>, tensor<i64>, tensor<20x40xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>, tensor<5x1x10xf32>, tensor<5x1x1xf32>, tensor<5x1x1xf32>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>>) -> tensor<40xf32> |
| %4 = "xla_hlo.get_tuple_element"(%arg0) {index = 3 : i32} : (tuple<tensor<i64>, tensor<i64>, tensor<40xf32>, tensor<i64>, tensor<20x40xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>, tensor<5x1x10xf32>, tensor<5x1x1xf32>, tensor<5x1x1xf32>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>>) -> tensor<i64> |
| %5 = "xla_hlo.get_tuple_element"(%arg0) {index = 4 : i32} : (tuple<tensor<i64>, tensor<i64>, tensor<40xf32>, tensor<i64>, tensor<20x40xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>, tensor<5x1x10xf32>, tensor<5x1x1xf32>, tensor<5x1x1xf32>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>>) -> tensor<20x40xf32> |
| %6 = "xla_hlo.get_tuple_element"(%arg0) {index = 5 : i32} : (tuple<tensor<i64>, tensor<i64>, tensor<40xf32>, tensor<i64>, tensor<20x40xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>, tensor<5x1x10xf32>, tensor<5x1x1xf32>, tensor<5x1x1xf32>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>>) -> tensor<i64> |
| %7 = "xla_hlo.get_tuple_element"(%arg0) {index = 9 : i32} : (tuple<tensor<i64>, tensor<i64>, tensor<40xf32>, tensor<i64>, tensor<20x40xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>, tensor<5x1x10xf32>, tensor<5x1x1xf32>, tensor<5x1x1xf32>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>>) -> tensor<5x1x1xf32> |
| %8 = "xla_hlo.gather"(%7, %0) {collapsed_slice_dims = dense<0> : tensor<1xi64>, index_vector_dim = 0 : i64, offset_dims = dense<[0, 1]> : tensor<2xi64>, slice_sizes = dense<1> : tensor<3xi64>, start_index_map = dense<0> : tensor<1xi64>} : (tensor<5x1x1xf32>, tensor<i64>) -> tensor<1x1xf32> |
| %9 = "xla_hlo.reshape"(%8) : (tensor<1x1xf32>) -> tensor<1xf32> |
| %10 = "xla_hlo.broadcast_in_dim"(%9) {broadcast_dimensions = dense<0> : tensor<1xi64>} : (tensor<1xf32>) -> tensor<1x10xf32> |
| %cst_0 = constant dense<1.0> : tensor<f32> |
| %11 = "xla_hlo.broadcast_in_dim"(%cst_0) : (tensor<f32>) -> tensor<1x10xf32> |
| %12 = xla_hlo.mul %10, %11 : tensor<1x10xf32> |
| %cst_1 = constant dense<0.0> : tensor<f32> |
| %13 = "xla_hlo.broadcast_in_dim"(%cst_1) : (tensor<f32>) -> tensor<1x10xf32> |
| %14 = "xla_hlo.compare"(%12, %13) {comparison_direction = "GT"} : (tensor<1x10xf32>, tensor<1x10xf32>) -> tensor<1x10xi1> |
| %15 = "xla_hlo.get_tuple_element"(%arg0) {index = 6 : i32} : (tuple<tensor<i64>, tensor<i64>, tensor<40xf32>, tensor<i64>, tensor<20x40xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>, tensor<5x1x10xf32>, tensor<5x1x1xf32>, tensor<5x1x1xf32>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>>) -> tensor<1x10xf32> |
| %cst_2 = constant dense<5.0e-01> : tensor<f32> |
| %16 = "xla_hlo.broadcast_in_dim"(%cst_2) : (tensor<f32>) -> tensor<1x10xf32> |
| %17 = "xla_hlo.broadcast_in_dim"(%cst_2) : (tensor<f32>) -> tensor<1x10xf32> |
| %18 = "xla_hlo.broadcast_in_dim"(%cst_2) : (tensor<f32>) -> tensor<1x10xf32> |
| %19 = "xla_hlo.get_tuple_element"(%arg0) {index = 8 : i32} : (tuple<tensor<i64>, tensor<i64>, tensor<40xf32>, tensor<i64>, tensor<20x40xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>, tensor<5x1x10xf32>, tensor<5x1x1xf32>, tensor<5x1x1xf32>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>>) -> tensor<5x1x10xf32> |
| %20 = "xla_hlo.gather"(%19, %0) {collapsed_slice_dims = dense<0> : tensor<1xi64>, index_vector_dim = 0 : i64, offset_dims = dense<[0, 1]> : tensor<2xi64>, slice_sizes = dense<[1, 1, 10]> : tensor<3xi64>, start_index_map = dense<0> : tensor<1xi64>} : (tensor<5x1x10xf32>, tensor<i64>) -> tensor<1x10xf32> |
| %21 = "xla_hlo.get_tuple_element"(%arg0) {index = 7 : i32} : (tuple<tensor<i64>, tensor<i64>, tensor<40xf32>, tensor<i64>, tensor<20x40xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>, tensor<5x1x10xf32>, tensor<5x1x1xf32>, tensor<5x1x1xf32>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>>) -> tensor<1x10xf32> |
| %22 = "xla_hlo.concatenate"(%20, %21) {dimension = 1 : i64} : (tensor<1x10xf32>, tensor<1x10xf32>) -> tensor<1x20xf32> |
| %23 = "xla_hlo.dot"(%22, %5) {precision_config = ["DEFAULT", "DEFAULT"]} : (tensor<1x20xf32>, tensor<20x40xf32>) -> tensor<1x40xf32> |
| %24 = "xla_hlo.transpose"(%23) {permutation = dense<[0, 1]> : tensor<2xi64>} : (tensor<1x40xf32>) -> tensor<1x40xf32> |
| %25 = "xla_hlo.reshape"(%3) : (tensor<40xf32>) -> tensor<1x40xf32> |
| %26 = xla_hlo.add %24, %25 : tensor<1x40xf32> |
| %27 = "xla_hlo.slice"(%26) {limit_indices = dense<[1, 30]> : tensor<2xi64>, start_indices = dense<[0, 20]> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<1x40xf32>) -> tensor<1x10xf32> |
| %28 = xla_hlo.mul %18, %27 : tensor<1x10xf32> |
| %29 = "xla_hlo.tanh"(%28) : (tensor<1x10xf32>) -> tensor<1x10xf32> |
| %30 = xla_hlo.mul %17, %29 : tensor<1x10xf32> |
| %31 = xla_hlo.add %16, %30 : tensor<1x10xf32> |
| %32 = xla_hlo.mul %31, %15 : tensor<1x10xf32> |
| %cst_3 = constant dense<5.0e-01> : tensor<f32> |
| %33 = "xla_hlo.broadcast_in_dim"(%cst_3) : (tensor<f32>) -> tensor<1x10xf32> |
| %34 = "xla_hlo.broadcast_in_dim"(%cst_3) : (tensor<f32>) -> tensor<1x10xf32> |
| %35 = "xla_hlo.broadcast_in_dim"(%cst_3) : (tensor<f32>) -> tensor<1x10xf32> |
| %36 = "xla_hlo.slice"(%26) {limit_indices = dense<[1, 20]> : tensor<2xi64>, start_indices = dense<[0, 10]> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<1x40xf32>) -> tensor<1x10xf32> |
| %37 = xla_hlo.mul %35, %36 : tensor<1x10xf32> |
| %38 = "xla_hlo.tanh"(%37) : (tensor<1x10xf32>) -> tensor<1x10xf32> |
| %39 = xla_hlo.mul %34, %38 : tensor<1x10xf32> |
| %40 = xla_hlo.add %33, %39 : tensor<1x10xf32> |
| %41 = "xla_hlo.slice"(%26) {limit_indices = dense<[1, 10]> : tensor<2xi64>, start_indices = dense<0> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<1x40xf32>) -> tensor<1x10xf32> |
| %42 = "xla_hlo.tanh"(%41) : (tensor<1x10xf32>) -> tensor<1x10xf32> |
| %43 = xla_hlo.mul %40, %42 : tensor<1x10xf32> |
| %44 = xla_hlo.add %32, %43 : tensor<1x10xf32> |
| %cst_4 = constant dense<1.0e+01> : tensor<f32> |
| %45 = "xla_hlo.broadcast_in_dim"(%cst_4) : (tensor<f32>) -> tensor<1x10xf32> |
| %46 = xla_hlo.min %44, %45 : tensor<1x10xf32> |
| %cst_5 = constant dense<-1.0e+01> : tensor<f32> |
| %47 = "xla_hlo.broadcast_in_dim"(%cst_5) : (tensor<f32>) -> tensor<1x10xf32> |
| %48 = xla_hlo.max %46, %47 : tensor<1x10xf32> |
| %49 = "xla_hlo.select"(%14, %15, %48) : (tensor<1x10xi1>, tensor<1x10xf32>, tensor<1x10xf32>) -> tensor<1x10xf32> |
| %50 = "xla_hlo.reshape"(%8) : (tensor<1x1xf32>) -> tensor<1xf32> |
| %51 = "xla_hlo.broadcast_in_dim"(%50) {broadcast_dimensions = dense<0> : tensor<1xi64>} : (tensor<1xf32>) -> tensor<1x10xf32> |
| %cst_6 = constant dense<1.0> : tensor<f32> |
| %52 = "xla_hlo.broadcast_in_dim"(%cst_6) : (tensor<f32>) -> tensor<1x10xf32> |
| %53 = xla_hlo.mul %51, %52 : tensor<1x10xf32> |
| %cst_7 = constant dense<0.0> : tensor<f32> |
| %54 = "xla_hlo.broadcast_in_dim"(%cst_7) : (tensor<f32>) -> tensor<1x10xf32> |
| %55 = "xla_hlo.compare"(%53, %54) {comparison_direction = "GT"} : (tensor<1x10xf32>, tensor<1x10xf32>) -> tensor<1x10xi1> |
| %cst_8 = constant dense<5.0e-01> : tensor<f32> |
| %56 = "xla_hlo.broadcast_in_dim"(%cst_8) : (tensor<f32>) -> tensor<1x10xf32> |
| %57 = "xla_hlo.broadcast_in_dim"(%cst_8) : (tensor<f32>) -> tensor<1x10xf32> |
| %58 = "xla_hlo.broadcast_in_dim"(%cst_8) : (tensor<f32>) -> tensor<1x10xf32> |
| %59 = "xla_hlo.slice"(%26) {limit_indices = dense<[1, 40]> : tensor<2xi64>, start_indices = dense<[0, 30]> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<1x40xf32>) -> tensor<1x10xf32> |
| %60 = xla_hlo.mul %58, %59 : tensor<1x10xf32> |
| %61 = "xla_hlo.tanh"(%60) : (tensor<1x10xf32>) -> tensor<1x10xf32> |
| %62 = xla_hlo.mul %57, %61 : tensor<1x10xf32> |
| %63 = xla_hlo.add %56, %62 : tensor<1x10xf32> |
| %64 = "xla_hlo.tanh"(%48) : (tensor<1x10xf32>) -> tensor<1x10xf32> |
| %65 = xla_hlo.mul %63, %64 : tensor<1x10xf32> |
| %66 = "xla_hlo.select"(%55, %21, %65) : (tensor<1x10xi1>, tensor<1x10xf32>, tensor<1x10xf32>) -> tensor<1x10xf32> |
| %67 = "xla_hlo.get_tuple_element"(%arg0) {index = 10 : i32} : (tuple<tensor<i64>, tensor<i64>, tensor<40xf32>, tensor<i64>, tensor<20x40xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>, tensor<5x1x10xf32>, tensor<5x1x1xf32>, tensor<5x1x1xf32>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>>) -> tensor<5x1x1xf32> |
| %68 = "xla_hlo.get_tuple_element"(%arg0) {index = 11 : i32} : (tuple<tensor<i64>, tensor<i64>, tensor<40xf32>, tensor<i64>, tensor<20x40xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>, tensor<5x1x10xf32>, tensor<5x1x1xf32>, tensor<5x1x1xf32>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>>) -> tensor<5xi64> |
| %69 = "xla_hlo.reshape"(%6) : (tensor<i64>) -> tensor<1xi64> |
| %70 = "xla_hlo.slice"(%69) {limit_indices = dense<1> : tensor<1xi64>, start_indices = dense<0> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<1xi64>) -> tensor<1xi64> |
| %71 = "xla_hlo.reshape"(%0) : (tensor<i64>) -> tensor<1xi64> |
| %72 = "xla_hlo.concatenate"(%71) {dimension = 0 : i64} : (tensor<1xi64>) -> tensor<1xi64> |
| %73 = "xla_hlo.convert"(%72) : (tensor<1xi64>) -> tensor<1xi32> |
| %74 = "xla_hlo.slice"(%73) {limit_indices = dense<1> : tensor<1xi64>, start_indices = dense<0> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<1xi32>) -> tensor<1xi32> |
| %75 = "xla_hlo.reshape"(%74) : (tensor<1xi32>) -> tensor<i32> |
| %76 = "xla_hlo.dynamic-update-slice"(%68, %70, %75) : (tensor<5xi64>, tensor<1xi64>, tensor<i32>) -> tensor<5xi64> |
| %77 = "xla_hlo.get_tuple_element"(%arg0) {index = 12 : i32} : (tuple<tensor<i64>, tensor<i64>, tensor<40xf32>, tensor<i64>, tensor<20x40xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>, tensor<5x1x10xf32>, tensor<5x1x1xf32>, tensor<5x1x1xf32>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>>) -> tensor<5x1x10xf32> |
| %78 = "xla_hlo.reshape"(%49) : (tensor<1x10xf32>) -> tensor<1x1x10xf32> |
| %79 = "xla_hlo.slice"(%78) {limit_indices = dense<[1, 1, 10]> : tensor<3xi64>, start_indices = dense<0> : tensor<3xi64>, strides = dense<1> : tensor<3xi64>} : (tensor<1x1x10xf32>) -> tensor<1x1x10xf32> |
| %80 = "xla_hlo.slice"(%73) {limit_indices = dense<1> : tensor<1xi64>, start_indices = dense<0> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<1xi32>) -> tensor<1xi32> |
| %81 = "xla_hlo.reshape"(%80) : (tensor<1xi32>) -> tensor<i32> |
| %cst_9 = constant dense<0> : tensor<i32> |
| %82 = "xla_hlo.dynamic-update-slice"(%77, %79, %81, %cst_9, %cst_9) : (tensor<5x1x10xf32>, tensor<1x1x10xf32>, tensor<i32>, tensor<i32>, tensor<i32>) -> tensor<5x1x10xf32> |
| %83 = "xla_hlo.get_tuple_element"(%arg0) {index = 13 : i32} : (tuple<tensor<i64>, tensor<i64>, tensor<40xf32>, tensor<i64>, tensor<20x40xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>, tensor<5x1x10xf32>, tensor<5x1x1xf32>, tensor<5x1x1xf32>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>>) -> tensor<5x1x10xf32> |
| %84 = "xla_hlo.reshape"(%66) : (tensor<1x10xf32>) -> tensor<1x1x10xf32> |
| %85 = "xla_hlo.slice"(%84) {limit_indices = dense<[1, 1, 10]> : tensor<3xi64>, start_indices = dense<0> : tensor<3xi64>, strides = dense<1> : tensor<3xi64>} : (tensor<1x1x10xf32>) -> tensor<1x1x10xf32> |
| %86 = "xla_hlo.slice"(%73) {limit_indices = dense<1> : tensor<1xi64>, start_indices = dense<0> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<1xi32>) -> tensor<1xi32> |
| %87 = "xla_hlo.reshape"(%86) : (tensor<1xi32>) -> tensor<i32> |
| %cst_10 = constant dense<0> : tensor<i32> |
| %88 = "xla_hlo.dynamic-update-slice"(%83, %85, %87, %cst_10, %cst_10) : (tensor<5x1x10xf32>, tensor<1x1x10xf32>, tensor<i32>, tensor<i32>, tensor<i32>) -> tensor<5x1x10xf32> |
| %89 = "xla_hlo.tuple"(%1, %2, %3, %4, %5, %6, %49, %66, %19, %7, %67, %76, %82, %88) : (tensor<i64>, tensor<i64>, tensor<40xf32>, tensor<i64>, tensor<20x40xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>, tensor<5x1x10xf32>, tensor<5x1x1xf32>, tensor<5x1x1xf32>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>) -> tuple<tensor<i64>, tensor<i64>, tensor<40xf32>, tensor<i64>, tensor<20x40xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>, tensor<5x1x10xf32>, tensor<5x1x1xf32>, tensor<5x1x1xf32>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>> |
| return %89 : tuple<tensor<i64>, tensor<i64>, tensor<40xf32>, tensor<i64>, tensor<20x40xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>, tensor<5x1x10xf32>, tensor<5x1x1xf32>, tensor<5x1x1xf32>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>> |
| } |
| func @ForwardLoopCond_gFAnjWGSoLs__.705(%arg0: tuple<tensor<i64>, tensor<i64>, tensor<40xf32>, tensor<i64>, tensor<20x40xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>, tensor<5x1x10xf32>, tensor<5x1x1xf32>, tensor<5x1x1xf32>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>>) -> tuple<tensor<i1>> { |
| %0 = "xla_hlo.get_tuple_element"(%arg0) {index = 0 : i32} : (tuple<tensor<i64>, tensor<i64>, tensor<40xf32>, tensor<i64>, tensor<20x40xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>, tensor<5x1x10xf32>, tensor<5x1x1xf32>, tensor<5x1x1xf32>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>>) -> tensor<i64> |
| %1 = "xla_hlo.get_tuple_element"(%arg0) {index = 1 : i32} : (tuple<tensor<i64>, tensor<i64>, tensor<40xf32>, tensor<i64>, tensor<20x40xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>, tensor<5x1x10xf32>, tensor<5x1x1xf32>, tensor<5x1x1xf32>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>>) -> tensor<i64> |
| %2 = "xla_hlo.compare"(%0, %1) {comparison_direction = "LT"} : (tensor<i64>, tensor<i64>) -> tensor<i1> |
| %3 = "xla_hlo.tuple"(%2) : (tensor<i1>) -> tuple<tensor<i1>> |
| return %3 : tuple<tensor<i1>> |
| } |
| func @cond_wrapper.723(%arg0: tuple<tensor<i64>, tensor<i64>, tensor<40xf32>, tensor<i64>, tensor<20x40xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>, tensor<5x1x10xf32>, tensor<5x1x1xf32>, tensor<5x1x1xf32>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>>) -> tensor<i1> { |
| %0 = call @ForwardLoopCond_gFAnjWGSoLs__.705(%arg0) : (tuple<tensor<i64>, tensor<i64>, tensor<40xf32>, tensor<i64>, tensor<20x40xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>, tensor<5x1x10xf32>, tensor<5x1x1xf32>, tensor<5x1x1xf32>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>>) -> tuple<tensor<i1>> |
| %1 = "xla_hlo.get_tuple_element"(%0) {index = 0 : i32} : (tuple<tensor<i1>>) -> tensor<i1> |
| return %1 : tensor<i1> |
| } |
| func @Forward_DIZIAkooG44__disable_call_shape_inference_true_.727(%arg0: tensor<1x10xf32>, %arg1: tensor<1x10xf32>, %arg2: tensor<5x1x10xf32>, %arg3: tensor<5x1x1xf32>, %arg4: tensor<5x1x1xf32>) -> tuple<tensor<i64>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>> { |
| %cst = constant dense<5> : tensor<i32> |
| %0 = "xla_hlo.convert"(%arg3) : (tensor<5x1x1xf32>) -> tensor<5x1x1xf32> |
| %cst_0 = constant dense<0x7F800000> : tensor<f32> |
| %1 = "xla_hlo.convert"(%cst_0) : (tensor<f32>) -> tensor<f32> |
| %2 = "xla_hlo.reduce"(%0, %1) ( { |
| ^bb0(%arg5: tensor<f32>, %arg6: tensor<f32>): // no predecessors |
| %42 = xla_hlo.min %arg5, %arg6 : tensor<f32> |
| "xla_hlo.return"(%42) : (tensor<f32>) -> () |
| }) {dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<5x1x1xf32>, tensor<f32>) -> tensor<5xf32> |
| %3 = "xla_hlo.convert"(%2) : (tensor<5xf32>) -> tensor<5xf32> |
| %cst_1 = constant dense<0.0> : tensor<f32> |
| %4 = "xla_hlo.broadcast_in_dim"(%cst_1) : (tensor<f32>) -> tensor<5xf32> |
| %5 = "xla_hlo.compare"(%3, %4) {comparison_direction = "EQ"} : (tensor<5xf32>, tensor<5xf32>) -> tensor<5xi1> |
| %6 = "xla_hlo.convert"(%5) : (tensor<5xi1>) -> tensor<5xi32> |
| %cst_2 = constant dense<[1, 2, 3, 4, 5]> : tensor<5xi32> |
| %7 = xla_hlo.mul %6, %cst_2 : tensor<5xi32> |
| %8 = "xla_hlo.convert"(%7) : (tensor<5xi32>) -> tensor<5xi32> |
| %cst_3 = constant dense<-2147483648> : tensor<i32> |
| %9 = "xla_hlo.convert"(%cst_3) : (tensor<i32>) -> tensor<i32> |
| %10 = "xla_hlo.reduce"(%8, %9) ( { |
| ^bb0(%arg5: tensor<i32>, %arg6: tensor<i32>): // no predecessors |
| %42 = xla_hlo.max %arg5, %arg6 : tensor<i32> |
| "xla_hlo.return"(%42) : (tensor<i32>) -> () |
| }) {dimensions = dense<0> : tensor<1xi64>} : (tensor<5xi32>, tensor<i32>) -> tensor<i32> |
| %11 = "xla_hlo.convert"(%10) : (tensor<i32>) -> tensor<i32> |
| %12 = xla_hlo.sub %cst, %11 : tensor<i32> |
| %cst_4 = constant dense<5> : tensor<i32> |
| %13 = "xla_hlo.compare"(%12, %cst_4) {comparison_direction = "EQ"} : (tensor<i32>, tensor<i32>) -> tensor<i1> |
| %cst_5 = constant dense<0> : tensor<i32> |
| %cst_6 = constant dense<5> : tensor<i32> |
| %14 = "xla_hlo.reverse"(%3) {dimensions = dense<0> : tensor<1xi64>} : (tensor<5xf32>) -> tensor<5xf32> |
| %cst_7 = constant dense<0.0> : tensor<f32> |
| %15 = "xla_hlo.broadcast_in_dim"(%cst_7) : (tensor<f32>) -> tensor<5xf32> |
| %16 = "xla_hlo.compare"(%14, %15) {comparison_direction = "EQ"} : (tensor<5xf32>, tensor<5xf32>) -> tensor<5xi1> |
| %17 = "xla_hlo.convert"(%16) : (tensor<5xi1>) -> tensor<5xi32> |
| %cst_8 = constant dense<[1, 2, 3, 4, 5]> : tensor<5xi32> |
| %18 = xla_hlo.mul %17, %cst_8 : tensor<5xi32> |
| %19 = "xla_hlo.convert"(%18) : (tensor<5xi32>) -> tensor<5xi32> |
| %cst_9 = constant dense<-2147483648> : tensor<i32> |
| %20 = "xla_hlo.convert"(%cst_9) : (tensor<i32>) -> tensor<i32> |
| %21 = "xla_hlo.reduce"(%19, %20) ( { |
| ^bb0(%arg5: tensor<i32>, %arg6: tensor<i32>): // no predecessors |
| %42 = xla_hlo.max %arg5, %arg6 : tensor<i32> |
| "xla_hlo.return"(%42) : (tensor<i32>) -> () |
| }) {dimensions = dense<0> : tensor<1xi64>} : (tensor<5xi32>, tensor<i32>) -> tensor<i32> |
| %22 = "xla_hlo.convert"(%21) : (tensor<i32>) -> tensor<i32> |
| %23 = xla_hlo.sub %cst_6, %22 : tensor<i32> |
| %24 = "xla_hlo.select"(%13, %cst_5, %23) : (tensor<i1>, tensor<i32>, tensor<i32>) -> tensor<i32> |
| %25 = "xla_hlo.convert"(%24) : (tensor<i32>) -> tensor<i64> |
| %cst_10 = constant dense<5> : tensor<i32> |
| %26 = xla_hlo.sub %cst_10, %12 : tensor<i32> |
| %27 = "xla_hlo.convert"(%26) : (tensor<i32>) -> tensor<i64> |
| %cst_11 = constant dense<0.0> : tensor<f32> |
| %28 = "xla_hlo.broadcast_in_dim"(%cst_11) : (tensor<f32>) -> tensor<40xf32> |
| %cst_12 = constant dense<0> : tensor<i64> |
| %cst_13 = constant dense<0.42> : tensor<20x40xf32> |
| %cst_14 = constant dense<-1356785880> : tensor<i64> |
| %cst_15 = constant dense<0> : tensor<i64> |
| %29 = "xla_hlo.broadcast_in_dim"(%cst_15) : (tensor<i64>) -> tensor<5xi64> |
| %cst_16 = constant dense<0.0> : tensor<f32> |
| %30 = "xla_hlo.broadcast_in_dim"(%cst_16) : (tensor<f32>) -> tensor<5x1x10xf32> |
| %cst_17 = constant dense<0.0> : tensor<f32> |
| %31 = "xla_hlo.broadcast_in_dim"(%cst_17) : (tensor<f32>) -> tensor<5x1x10xf32> |
| %32 = "xla_hlo.tuple"(%25, %27, %28, %cst_12, %cst_13, %cst_14, %arg0, %arg1, %arg2, %arg3, %arg4, %29, %30, %31) : (tensor<i64>, tensor<i64>, tensor<40xf32>, tensor<i64>, tensor<20x40xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>, tensor<5x1x10xf32>, tensor<5x1x1xf32>, tensor<5x1x1xf32>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>) -> tuple<tensor<i64>, tensor<i64>, tensor<40xf32>, tensor<i64>, tensor<20x40xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>, tensor<5x1x10xf32>, tensor<5x1x1xf32>, tensor<5x1x1xf32>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>> |
| %33 = "xla_hlo.while"(%32) {body = @ForwardLoopBody_o0Jnom3Cdxo__.597, cond = @cond_wrapper.723} : (tuple<tensor<i64>, tensor<i64>, tensor<40xf32>, tensor<i64>, tensor<20x40xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>, tensor<5x1x10xf32>, tensor<5x1x1xf32>, tensor<5x1x1xf32>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>>) -> tuple<tensor<i64>, tensor<i64>, tensor<40xf32>, tensor<i64>, tensor<20x40xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>, tensor<5x1x10xf32>, tensor<5x1x1xf32>, tensor<5x1x1xf32>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>> |
| %34 = "xla_hlo.get_tuple_element"(%33) {index = 0 : i32} : (tuple<tensor<i64>, tensor<i64>, tensor<40xf32>, tensor<i64>, tensor<20x40xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>, tensor<5x1x10xf32>, tensor<5x1x1xf32>, tensor<5x1x1xf32>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>>) -> tensor<i64> |
| %35 = "xla_hlo.get_tuple_element"(%33) {index = 11 : i32} : (tuple<tensor<i64>, tensor<i64>, tensor<40xf32>, tensor<i64>, tensor<20x40xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>, tensor<5x1x10xf32>, tensor<5x1x1xf32>, tensor<5x1x1xf32>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>>) -> tensor<5xi64> |
| %36 = "xla_hlo.get_tuple_element"(%33) {index = 12 : i32} : (tuple<tensor<i64>, tensor<i64>, tensor<40xf32>, tensor<i64>, tensor<20x40xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>, tensor<5x1x10xf32>, tensor<5x1x1xf32>, tensor<5x1x1xf32>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>>) -> tensor<5x1x10xf32> |
| %37 = "xla_hlo.get_tuple_element"(%33) {index = 13 : i32} : (tuple<tensor<i64>, tensor<i64>, tensor<40xf32>, tensor<i64>, tensor<20x40xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>, tensor<5x1x10xf32>, tensor<5x1x1xf32>, tensor<5x1x1xf32>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>>) -> tensor<5x1x10xf32> |
| %38 = "xla_hlo.get_tuple_element"(%33) {index = 5 : i32} : (tuple<tensor<i64>, tensor<i64>, tensor<40xf32>, tensor<i64>, tensor<20x40xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>, tensor<5x1x10xf32>, tensor<5x1x1xf32>, tensor<5x1x1xf32>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>>) -> tensor<i64> |
| %39 = "xla_hlo.get_tuple_element"(%33) {index = 6 : i32} : (tuple<tensor<i64>, tensor<i64>, tensor<40xf32>, tensor<i64>, tensor<20x40xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>, tensor<5x1x10xf32>, tensor<5x1x1xf32>, tensor<5x1x1xf32>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>>) -> tensor<1x10xf32> |
| %40 = "xla_hlo.get_tuple_element"(%33) {index = 7 : i32} : (tuple<tensor<i64>, tensor<i64>, tensor<40xf32>, tensor<i64>, tensor<20x40xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>, tensor<5x1x10xf32>, tensor<5x1x1xf32>, tensor<5x1x1xf32>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>>) -> tensor<1x10xf32> |
| %41 = "xla_hlo.tuple"(%34, %35, %36, %37, %38, %39, %40) : (tensor<i64>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>) -> tuple<tensor<i64>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>> |
| return %41 : tuple<tensor<i64>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>> |
| } |
| |
| // CHECK-LABEL: EXEC @main |
| func @main(%arg0: tensor<1x5xf32>, %arg1: tensor<1x5x2x2xf32>) -> tuple<tensor<5x1x10xf32>> |
| attributes { iree.module.export } { |
| %cst = constant dense<0.0> : tensor<f32> |
| %0 = "xla_hlo.broadcast_in_dim"(%cst) : (tensor<f32>) -> tensor<1x10xf32> |
| %cst_0 = constant dense<0.0> : tensor<f32> |
| %1 = "xla_hlo.broadcast_in_dim"(%cst_0) : (tensor<f32>) -> tensor<1x10xf32> |
| %cst_1 = constant dense<0.0> : tensor<f32> |
| %2 = "xla_hlo.broadcast_in_dim"(%cst_1) : (tensor<f32>) -> tensor<1x10xf32> |
| %cst_2 = constant dense<0.0> : tensor<f32> |
| %3 = "xla_hlo.broadcast_in_dim"(%cst_2) : (tensor<f32>) -> tensor<1x10xf32> |
| %cst_3 = constant dense<0.0> : tensor<f32> |
| %4 = "xla_hlo.broadcast_in_dim"(%cst_3) : (tensor<f32>) -> tensor<1x10xf32> |
| %cst_4 = constant dense<0.0> : tensor<f32> |
| %5 = "xla_hlo.broadcast_in_dim"(%cst_4) : (tensor<f32>) -> tensor<1x10xf32> |
| %6 = "xla_hlo.reshape"(%arg1) : (tensor<1x5x2x2xf32>) -> tensor<1x5x2x2xf32> |
| %7 = "xla_hlo.reshape"(%6) : (tensor<1x5x2x2xf32>) -> tensor<1x5x4xf32> |
| %cst_5 = constant dense<0.0> : tensor<f32> |
| %8 = "xla_hlo.pad"(%7, %cst_5) {edge_padding_high = dense<[0, 0, 60]> : tensor<3xi64>, edge_padding_low = dense<0> : tensor<3xi64>, interior_padding = dense<0> : tensor<3xi64>} : (tensor<1x5x4xf32>, tensor<f32>) -> tensor<1x5x64xf32> |
| %9 = "xla_hlo.transpose"(%8) {permutation = dense<[1, 0, 2]> : tensor<3xi64>} : (tensor<1x5x64xf32>) -> tensor<5x1x64xf32> |
| %10 = "xla_hlo.reshape"(%arg0) : (tensor<1x5xf32>) -> tensor<1x5xf32> |
| %11 = "xla_hlo.transpose"(%10) {permutation = dense<[1, 0]> : tensor<2xi64>} : (tensor<1x5xf32>) -> tensor<5x1xf32> |
| %12 = "xla_hlo.reshape"(%11) : (tensor<5x1xf32>) -> tensor<5x1x1xf32> |
| %cst_6 = constant dense<0.0> : tensor<f32> |
| %13 = "xla_hlo.broadcast_in_dim"(%cst_6) : (tensor<f32>) -> tensor<5x1x1xf32> |
| %14 = call @Forward_DIZIAkooG44__disable_call_shape_inference_true_.189(%4, %5, %9, %12, %13) : (tensor<1x10xf32>, tensor<1x10xf32>, tensor<5x1x64xf32>, tensor<5x1x1xf32>, tensor<5x1x1xf32>) -> tuple<tensor<i64>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>> |
| %15 = "xla_hlo.get_tuple_element"(%14) {index = 3 : i32} : (tuple<tensor<i64>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>>) -> tensor<5x1x10xf32> |
| %cst_7 = constant dense<0.0> : tensor<f32> |
| %16 = "xla_hlo.broadcast_in_dim"(%cst_7) : (tensor<f32>) -> tensor<5x1x1xf32> |
| %17 = call @Forward_DIZIAkooG44__disable_call_shape_inference_true_.458(%2, %3, %15, %12, %16) : (tensor<1x10xf32>, tensor<1x10xf32>, tensor<5x1x10xf32>, tensor<5x1x1xf32>, tensor<5x1x1xf32>) -> tuple<tensor<i64>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>> |
| %18 = "xla_hlo.get_tuple_element"(%17) {index = 3 : i32} : (tuple<tensor<i64>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>>) -> tensor<5x1x10xf32> |
| %cst_8 = constant dense<0.0> : tensor<f32> |
| %19 = "xla_hlo.broadcast_in_dim"(%cst_8) : (tensor<f32>) -> tensor<5x1x1xf32> |
| %20 = call @Forward_DIZIAkooG44__disable_call_shape_inference_true_.727(%0, %1, %18, %12, %19) : (tensor<1x10xf32>, tensor<1x10xf32>, tensor<5x1x10xf32>, tensor<5x1x1xf32>, tensor<5x1x1xf32>) -> tuple<tensor<i64>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>> |
| %21 = "xla_hlo.get_tuple_element"(%20) {index = 3 : i32} : (tuple<tensor<i64>, tensor<5xi64>, tensor<5x1x10xf32>, tensor<5x1x10xf32>, tensor<i64>, tensor<1x10xf32>, tensor<1x10xf32>>) -> tensor<5x1x10xf32> |
| %22 = "xla_hlo.copy"(%21) : (tensor<5x1x10xf32>) -> tensor<5x1x10xf32> |
| %23 = "xla_hlo.reshape"(%22) : (tensor<5x1x10xf32>) -> tensor<5x1x10xf32> |
| %24 = "xla_hlo.tuple"(%23) : (tensor<5x1x10xf32>) -> tuple<tensor<5x1x10xf32>> |
| return %24 : tuple<tensor<5x1x10xf32>> |
| } |
| |
| // CHECK: 5x1x10xf32= |
| // CHECK-SAME: [ |
| // CHECK-SAME: [0.7{{[0-9]+}} 0.7{{[0-9]+}} 0.7{{[0-9]+}} 0.7{{[0-9]+}} 0.7{{[0-9]+}} 0.7{{[0-9]+}} 0.7{{[0-9]+}} 0.7{{[0-9]+}} 0.7{{[0-9]+}} 0.7{{[0-9]+}}] |
| // CHECK-SAME: ][ |
| // CHECK-SAME: [0.7{{[0-9]+}} 0.7{{[0-9]+}} 0.7{{[0-9]+}} 0.7{{[0-9]+}} 0.7{{[0-9]+}} 0.7{{[0-9]+}} 0.7{{[0-9]+}} 0.7{{[0-9]+}} 0.7{{[0-9]+}} 0.7{{[0-9]+}}] |
| // CHECK-SAME: ][ |
| // CHECK-SAME: [0.9{{[0-9]+}} 0.9{{[0-9]+}} 0.9{{[0-9]+}} 0.9{{[0-9]+}} 0.9{{[0-9]+}} 0.9{{[0-9]+}} 0.9{{[0-9]+}} 0.9{{[0-9]+}} 0.9{{[0-9]+}} 0.9{{[0-9]+}}] |
| // CHECK-SAME: ][ |
| // CHECK-SAME: [0 0 0 0 0 0 0 0 0 0] |
| // CHECK-SAME: ][ |
| // CHECK-SAME: [0 0 0 0 0 0 0 0 0 0] |
| // CHECK-SAME: ] |
| |