%0 = "tf.VarHandleOp"() {container = "c", shared_name = "v"} : () -> tensor<*x!tf_type.resource<tensor<2x8xi32>>>
    %1 = "tf.VarHandleOp"() {container = "c", shared_name = "w"} : () -> tensor<*x!tf_type.resource<tensor<2xi32>>>
    %2 = "tf.Const"() {value = dense<1> : tensor<i32>} : () -> tensor<i32>
    %3 = "tf.Const"() {value = dense<1> : tensor<i32>} : () -> tensor<i32>

// CHECK-LABEL: func @main
func.func @main(%arg0 : tensor<10xf32>, %arg1 : tensor<10xf32>) -> tensor<10x10xf32> {
// CHECK-SAME: (%[[ARG0:.*]]: tensor<10xf32>, %[[ARG1:.*]]: tensor<10xf32>)
// CHECK-NEXT:    %[[ADD:.*]] = "tf.Add"(%[[ARG0]], %[[ARG1]]) : (tensor<10xf32>, tensor<10xf32>) -> tensor<10xf32>
// CHECK-NEXT:    %[[MAGIC:.*]] = "magic.op"(%[[ADD]], %[[ADD]]) : (tensor<10xf32>, tensor<10xf32>) -> tensor<10x10xf32>
// CHECK-NEXT:    return %[[MAGIC]]

  attr {
    key: "dtype"
    value {
      type: DT_INT32
    }
  }
}
versions {
  producer: 27
}

# CHECK-LABEL: func @main
# CHECK-SAME:  (%[[ARG_0:[a-z0-9]+]]: tensor<10xi32>, %[[ARG_1:[a-z0-9]+]]: tensor<10xi32>) -> tensor<*xi32>
# CHECK-SAME:  control_outputs = ""
# CHECK-SAME:  inputs = "input0,input1"
# CHECK-SAME:  outputs = "Sub"

  // CHECK-DAG: %[[RHS_1:.*]] = mhlo.convert %arg1 : tensor<2xi32>
  // CHECK-DAG: %[[RHS_2:.*]] = mhlo.bitcast_convert %[[RHS_1]] : (tensor<2xi32>) -> tensor<2x!quant.uniform<i32:f32, 2.000000e+00:4>>
  %input1_qint = "tf.Cast"(%input1) {Truncate = false} : (tensor<3x2xi32>) -> tensor<3x2x!tf_type.qint32>
  %input2_qint = "tf.Cast"(%input2) {Truncate = false} : (tensor<2xi32>) -> tensor<2x!tf_type.qint32>

  func.return %0, %1: tensor<256x256xi32>, tensor<128xi32>
}

func.func @while_body(%arg0: tensor<*xi32>) -> (tensor<*xi32>) {
  %0 = "tf.Rank"(%arg0) : (tensor<*xi32>) -> tensor<i32>
  %1 = "tf.Add"(%0, %arg0): (tensor<i32>, tensor<*xi32>) -> tensor<*xi32>
  func.return %1: tensor<*xi32>
}

    } : (tensor<10xf32>, tensor<10xf32>, tensor<6xi32>) -> (tensor<10xf32>, tensor<6xi32>)
    // CHECK: return %[[RESULT_CAST]], %[[RESULT]]#1 : tensor<10xf32>, tensor<6xi32>
    func.return %0#0, %0#1 : tensor<10xf32>, tensor<6xi32>
  }

  // CHECK-LABEL: func private @main_0
  // CHECK-SAME: %[[ARG_0:.*]]: tensor<10xbf16>, %[[ARG_1:.*]]: tensor<10xbf16>, %[[ARG_2:.*]]: tensor<6xi32>
  func.func private @main_0(

    %1 = "tf.ReadVariableOp"(%0) : (tensor<!tf_type.resource<tensor<10xf32>>>) -> tensor<10xf32>
    %2 = "tf.ReadVariableOp"(%0) : (tensor<!tf_type.resource<tensor<10xf32>>>) -> tensor<10xf32>
    %3 = "tf.Mul"(%1, %2) : (tensor<10xf32>, tensor<10xf32>) -> tensor<10xf32>
    return %3 : tensor<10xf32>
  }
}

// -----

// CHECK-LABEL: module

func.func @xla_gather_known_output_shape(%arg0: tensor<5xi32>, %arg1: tensor<1xi32>, %arg2: tensor<1xi32>) -> tensor<i32> {
  // dimension_numbers: {
  //   collapsed_slice_dims: 0
  //   start_index_map: 0
  // }
  %0 = "tf.XlaGather"(%arg0, %arg1, %arg2) {device = "", dimension_numbers = "\12\01\00\1A\01\00", indices_are_sorted = true} : (tensor<5xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<i32>

// CHECK-DAG: %[[CST_5:.*]] = "tf.Const"() <{value = dense<-19> : tensor<i32>}> : () -> tensor<i32>
// CHECK-DAG: %[[CST_6:.*]] = "tf.Const"() <{value = dense<0> : tensor<2xi32>}> : () -> tensor<2xi32>
// CHECK-NEXT: %[[V0:.*]] = "tf.PartitionedCall"(%[[ARG1]], %[[CST_1]], %[[CST_0]]) <{config = "", config_proto = "", executor_type = "", f = @quantize_i8}> : (tensor<2x3x3x2xf32>, tensor<f32>, tensor<i32>) -> tensor<2x3x3x2xi8>

// CHECK:           [[VAL_7:%.*]] = "tfl.slice"([[VAL_1]], [[VAL_3]], [[VAL_5]]) : (tensor<1x8x8x1024xf32>, tensor<4xi32>, tensor<4xi32>) -> tensor<1x8x8x256xf32>
// CHECK:           [[VAL_8:%.*]] = "tfl.slice"([[VAL_1]], [[VAL_4]], [[VAL_5]]) : (tensor<1x8x8x1024xf32>, tensor<4xi32>, tensor<4xi32>) -> tensor<1x8x8x256xf32>