%1 = "tf.TensorListGetItem"(%0, %arg2, %arg1) {device = "/job:localhost/replica:0/task:0/device:GPU:0"} : (tensor<!tf_type.variant<tensor<1x32xf32>>>, tensor<i32>, tensor<2xi32>) -> tensor<1x32xf32>
  %2 = "tf.TensorListGetItem"(%0, %arg3, %arg1) {device = "/job:localhost/replica:0/task:0/device:GPU:0"} : (tensor<!tf_type.variant<tensor<1x32xf32>>>, tensor<i32>, tensor<2xi32>) -> tensor<1x32xf32>

  %2 = "tf.BatchToSpaceND"(%arg0, %0, %1) {device = ""} : (tensor<3x5x2xf32>, tensor<1xi32>, tensor<1x2xi32>) -> tensor<1x8x2xf32>

  // CHECK: return [[VAL8]] : tensor<1x8x2xf32>
  func.return %2 : tensor<1x8x2xf32>
}

func.func @fake_quant_with_min_max_args(%arg0 : tensor<?x?xf32>) -> tensor<?x?xf32> {

  // tf.XlaReduceScatter legalization
  //===--------------------------------------------------------------------===//
  // CHECK-LABEL: func @xla_reduce_scatter
  func.func @xla_reduce_scatter(%arg0: tensor<128x128xf32>) -> tensor<64x128xf32> {
      %cst = "tf.Const"() {value = dense<0> : tensor<i32>} : () -> tensor<i32>
      %cst_0 = "tf.Const"() {value = dense<[[0, 4], [1, 5], [2, 6], [3, 7]]> : tensor<4x2xi32>} : () -> tensor<4x2xi32>

// 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

  %tl0 = "tf.EmptyTensorList"(%element_shape, %max_size) : (tensor<1xi32>, tensor<i32>) -> tensor<!tf_type.variant<tensor<16xf32>>>
  // CHECK-NOT: tf.TensorListReserve
  // CHECK: tf.BroadcastTo
  // CHECK-SAME: tensor<10x32xf32>
  %tl1 = "tf.TensorListReserve"(%arg0, %max_size) : (tensor<*xi32>, tensor<i32>) -> tensor<!tf_type.variant<tensor<32xf32>>>
  func.return
}

// -----

// dimensions into a single dimension. For example,
//
//   %shape = arith.constant dense<[1, 128, 64]> : tensor<3xi32>
//   %reshape = tfl.reshape(%input, %shape) // %input: tensor<128x64xf32>
//   %fc = tfl.fully_connected(%reshape, %filter, %bias)
//           {keep_num_dims = false, weights_format = "DEFAULT"}
//
// can be canonicalized to
//