func.return %4#0, %4#1 : tensor<4xf32>, tensor<4xi32>

  // CHECK-DAG:  %0 = mhlo.constant dense<0xFF800000> : tensor<f32>
  // CHECK-DAG:  %1 = mhlo.constant dense<0> : tensor<i32>
  // CHECK:  %2 = mhlo.constant dense<{{\[\[}}0, 0, 0, 0], [1, 1, 1, 1], [2, 2, 2, 2], [3, 3, 3, 3]]> : tensor<4x4xi32>

func.func @dequantize_pushdown(%arg0: tensor<1000x1000x!quant.uniform<i8:f32, 7.812500e-03>>, %arg1: tensor<1x1xi32>) -> tensor<1x1x1000xf32> {
  %0 = "tfl.dequantize"(%arg0) : (tensor<1000x1000x!quant.uniform<i8:f32, 7.812500e-03>>) -> tensor<1000x1000xf32>
  %1 = "tfl.gather"(%0, %arg1) {axis = 0 : i32, batch_dims = 0 : i32}: (tensor<1000x1000xf32>, tensor<1x1xi32>) -> tensor<1x1x1000xf32>
  func.return %1 : tensor<1x1x1000xf32>

// CHECK-NEXT: tfl.gather

func.func @testSubOfZero(%arg0: tensor<?x1xf32>, %arg1: tensor<4x1xf32>) -> (tensor<?x1xf32>, tensor<4x1xf32>) {
  %0 = "tf.Const"() {value = dense<0.0> : tensor<f32>} : () -> tensor<f32>
  %1 = "tf.Sub"(%arg0, %0) : (tensor<?x1xf32>, tensor<f32>) -> tensor<?x1xf32>
  %2 = "tf.Sub"(%arg1, %0) : (tensor<4x1xf32>, tensor<f32>) -> tensor<4x1xf32>
  func.return %1, %2: tensor<?x1xf32>, tensor<4x1xf32>

// CHECK: return %arg0, %arg1
}

func.func @decompose_scalar_updates(%arg0: tensor<!tf_type.resource<tensor<4x3xi32>>>, %arg1: tensor<2xi32>, %arg2: tensor<i32>) {
  "tf_device.cluster"() ({
    // CHECK-NOT: ResourceScatterUpdate
    // CHECK-NOT: BroadcastTo
    // CHECK: TensorScatterUpdate
    "tf.ResourceScatterUpdate"(%arg0, %arg1, %arg2) : (tensor<!tf_type.resource<tensor<4x3xi32>>>, tensor<2xi32>, tensor<i32>) -> ()
    tf_device.return
  }) : () -> ()

// -----

func.func @constant_tensor_invalid_3(%arg0: vector<1xi32>) -> tensor<1x1xi32> {
    // expected-error@+1 {{input and output should have same shape and element type}}
  %0 = "tfr.constant_tensor"(%arg0) : (vector<1xi32>) -> tensor<1x1xi32>
  func.return %0 : tensor<1x1xi32>
}

// -----

func.func @constant_tensor_invalid_4(%arg0: i32) -> tensor<1x1xi32> {
    // expected-error@+1 {{input can not be converted to an output tensor}}

    %22 = "tf.Cast"(%19) {Truncate = false} : (tensor<2x?x?xi8>) -> tensor<2x?x?xi32>
    %23 = "tf.Sum"(%22, %cst_3) {keep_dims = true} : (tensor<2x?x?xi32>, tensor<1xi64>) -> tensor<2x?x1xi32>
    %24 = "tf.Mul"(%23, %cst_4) : (tensor<2x?x1xi32>, tensor<i32>) -> tensor<2x?x1xi32>
    %25 = "tf.Cast"(%20) {Truncate = false} : (tensor<2x?x?xi8>) -> tensor<2x?x?xi32>

      // CHECK:          "mhlo.reduce_scatter"(%arg0)
      // CHECK{LITERAL}: replica_groups = dense<[[0, 4], [1, 5], [2, 6], [3, 7]]>
      // CHECK-SAME:     scatter_dimension = 0
      //
      %1 = "tf.XlaReduceScatter"(%arg0, %cst_0, %cst) {reduce_op = "Add"} : (tensor<128x128xf32>, tensor<4x2xi32>, tensor<i32>) -> tensor<64x128xf32>

  %cst = "tfl.pseudo_const"() {value = dense<[[0], [0], [1], [1], [2], [2], [3], [3], [4], [4], [5], [5], [6], [6], [7], [7]]> : tensor<16x1xi32>} : () -> tensor<16x1xi32>
  %gather_nd_first = "tfl.gather_nd"(%arg0, %cst) : (tensor<8x1x8x1280xf32>, tensor<16x1xi32>) -> tensor<16x1x8x1280xf32>
  %transpose_first_perm = "tfl.pseudo_const"() {value = dense<[2, 1, 0, 3]> : tensor<4xi32>} : () -> tensor<4xi32>

func.func @scatter_nd_add(%arg0: tensor<7xi64>, %arg1: tensor<1x1xi32>, %arg2: tensor<1xi64>) -> tensor<7xi64> {
  %0 = "tf.TensorScatterAdd"(%arg0, %arg1, %arg2) : (tensor<7xi64>, tensor<1x1xi32>, tensor<1xi64>) -> tensor<7xi64>
  func.return %0 : tensor<7xi64>

  // CHECK-LABEL: scatter_nd_add
  // CHECK:  %[[GATHER:.*]] = "tf.GatherNd"(%arg0, %arg1) <{bad_indices_policy = ""}> : (tensor<7xi64>, tensor<1x1xi32>) -> tensor<1xi64>

  %cst = "tf.Const"() {device = "", value = dense<2> : tensor<i32>} : () -> tensor<i32>
  // idx 1 should be a compile time constant
  %0 = "tf.ExpandDims"(%arg0, %cst) {device = ""} : (tensor<?x2xi32>, tensor<i32>) -> tensor<?x2x1xi32>
  %1 = "tf.AddV2"(%cst, %cst) {device = ""} : (tensor<i32>, tensor<i32>) -> tensor<i32>

  return %0 : tensor<?x2x1xi32>
}