}

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>

  %1 = "tfl.shape"(%conv_s) : (tensor<?x26x26x26x16xf32>) -> tensor<5xi64>
  %2 = "tfl.broadcast_args"(%1, %cst) : (tensor<5xi64>, tensor<1xi64>) -> tensor<5xi64>
  %broad1 = "tfl.broadcast_to"(%conv_s, %2) : (tensor<?x26x26x26x16xf32>, tensor<5xi64>) -> tensor<?x26x26x26x16xf32>
  %broad2 = "tfl.broadcast_to"(%b, %2) : (tensor<16xf32>, tensor<5xi64>) -> tensor<?x26x26x26x16xf32>

  // CHECK-DAG: %[[X_NORMED:.+]] = mhlo.multiply %[[X]], %[[MULTIPLIER_BCAST]] : tensor<4x256xf32>
  // CHECK-DAG: %[[RHS_BCAST:.+]] = "mhlo.broadcast_in_dim"(%[[RHS]]) <{broadcast_dimensions = dense<1> : tensor<1xi64>}> : (tensor<256xf32>) -> tensor<4x256xf32>

func.func @Int64SliceBeginSize(%arg0: tensor<4x128x32xf32>) -> tensor<1x128x32xf32> {
  %0 = "tfl.pseudo_const"() {value = dense<0> : tensor<3xi64>} : () -> tensor<3xi64>
  %1 = "tfl.pseudo_const"() {value = dense<[1, 128, 32]> : tensor<3xi64>} : () -> tensor<3xi64>
  %2 = "tfl.slice"(%arg0, %0, %1) : (tensor<4x128x32xf32>, tensor<3xi64>, tensor<3xi64>) -> tensor<1x128x32xf32>
  func.return %2 : tensor<1x128x32xf32>

    %shape_64 = "tf.TensorListElementShape"(%tl_1) : (tensor<!tf_type.variant<tensor<?x1xf32>>>) -> tensor<?xi64>
    // CHECK: %[[CAST:.*]] = "tf.Cast"(%[[ELEMENT_SHAPE]]){{.*}}: (tensor<2xi32>) -> tensor<2xi64>
    // CHECK: "tf._SomeOtherOp"(%[[ELEMENT_SHAPE]], %[[CAST]])
    "tf._SomeOtherOp"(%shape_32, %shape_64) : (tensor<?xi32>, tensor<?xi64>) -> ()
    func.return
  }

  // CHECK-LABEL: refine_pop_back_results_from_operands

  %2 = mhlo.constant dense<0.000000e+00> : tensor<f32>

// CHECK: return %[[GATHER]] : tensor<1x300x10xf32>

// -----

module {
  func.func @propagate_xlagather(%arg0: tensor<10x2xi32>) -> tensor<1x300x10xf32> {
    %cst = "tf.Const"() {value = dense<[1, 1, 300]> : tensor<3xi64>} : () -> tensor<3xi64>

  // CHECK-NEXT:  return [[FIND]] : tensor<*xi64>
  func.return %1 : tensor<*xi64>
}

// -----

// Test for case with int64 -> string hashtable.
func.func @hashtable_int64_to_string(%arg0: tensor<i64>) -> tensor<*x!tf_type.string> {
  %cst = arith.constant dense<"f"> : tensor<!tf_type.string>

  // CHECK: %[[CST:.*]] = "tf.Const"() <{value = dense<[0, 2, 3, 1]> : tensor<4xi64>}>
  // CHECK: %[[R0:.*]] = "tf.Transpose"(%[[ARG0]], %[[CST]])
  // CHECK: %[[R1:.*]] = "tf.BiasAdd"(%[[R0]], %[[ARG1]]) <{data_format = "NHWC"}> {device = ""}
  // CHECK: %[[CST_0:.*]] = "tf.Const"() <{value = dense<[0, 3, 1, 2]> : tensor<4xi64>}>
  // CHECK: "tf.Transpose"(%[[R1]], %[[CST_0]])

  %2 = "mhlo.dynamic_broadcast_in_dim"(%cst_0, %1) <{broadcast_dimensions = dense<3> : tensor<1xi64>}> : (tensor<2xf32>, tensor<4xindex>) -> tensor<?x256x256x2xf32>
  %3 = mhlo.multiply %0, %2 : tensor<?x256x256x2xf32>
  %4 = "mhlo.dynamic_broadcast_in_dim"(%cst_1, %1) <{broadcast_dimensions = dense<3> : tensor<1xi64>}> : (tensor<2xf32>, tensor<4xindex>) -> tensor<?x256x256x2xf32>
  %5 = mhlo.add %3, %4 : tensor<?x256x256x2xf32>