tf_device.return %2 : tensor<1x8x2xf32>
    // CHECK: () -> tensor<1x8x2xf32>
    }) {device = "/device:CPU:0"} : () -> tensor<*xf32>
    // CHECK: "tf.Cast"(%{{.*}}) <{Truncate = false}> : (tensor<1x8x2xf32>) -> tensor<*xf32>
    // CHECK: (tensor<i32>, tensor<1x8x2xf32>) -> (tensor<1x8x1xf32>, tensor<1x8x1xf32>)

// CHECK:           [[VAL_12:%.*]]:4 = "tf.SplitV"([[VAL_7]], [[VAL_10]], [[VAL_11]]) : (tensor<40x8xf32>, tensor<4xi32>, tensor<i32>) -> (tensor<10x8xf32>, tensor<10x8xf32>, tensor<10x8xf32>, tensor<10x8xf32>)
// CHECK-DAG:       [[VAL_13:%.*]] = "tf.Const"() <{value = dense<10> : tensor<4xi32>}> : () -> tensor<4xi32>

// CHECK: return %[[RES]]
}

func.func @scatterNdHigherRankIndices(%arg0: tensor<4x2x2xi32>, %arg1: tensor<4x2x3xf32>, %arg2: tensor<3xi32>) -> tensor<10x2x3xf32> {
  %0 = "tf.ScatterNd"(%arg0, %arg1, %arg2) : (tensor<4x2x2xi32>, tensor<4x2x3xf32>, tensor<3xi32>) -> tensor<10x2x3xf32>
  func.return %0 : tensor<10x2x3xf32>

// CHECK-LABEL:scatterNdHigherRankIndices

  %push = "tf.XlaDynamicUpdateSlice"(%tl, %arg0_reshaped, %start_indices) : (tensor<10x8x4xf32>, tensor<1x8x4xf32>, tensor<3xi32>) -> tensor<10x8x4xf32>
  %one = "tf.Const"() {value = dense<1> : tensor<1xi32>} : () -> tensor<1xi32>
  %next_index_in_list = "tf.AddV2"(%index_in_list, %one) : (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi32>
  return

// result shapes.
// CHECK-LABEL: func @testShapeInvariantWhile
func.func @testShapeInvariantWhile(%arg0: tensor<1x2x3xf32>) -> tensor<1x8x3xf32> {
  %0 = "tf.While"(%arg0) {cond = @cond, body = @body, is_stateless = false, shape_invariant} : (tensor<1x2x3xf32>) -> tensor<1x8x3xf32>
  func.return %0 : tensor<1x8x3xf32>
}

// -----
// WhileRegion tests

// Simple While region
// CHECK-LABEL: testValidWhileRegion

      // were:
      //   tensor<!tf_type.variant<tensor<?x8xf32>>>
      // and:
      //   tensor<!tf_type.variant<tensor<10x8xf32>>>
      // we'll try here to refine tensor<?x8xf32> with tensor<10x8xf32>.
      auto refined_subtype = mlir::cast<TensorType>(
          TypeMeet(lhs_element_type_with_subtype.GetSubtypes().front(),