// CHECK: return %[[v4]] : tensor<?x?x8x?xf32>
}

func.func @einsum_ellipsis(%arg0: tensor<1x512x128xf32>, %arg1: tensor<128x256xf32>) -> tensor<1x512x256xf32> {
  %0 = "tf.Einsum"(%arg0, %arg1) {device = "", equation = "...x,xy->...y"} : (tensor<1x512x128xf32>, tensor<128x256xf32>) -> tensor<1x512x256xf32>
  func.return %0 : tensor<1x512x256xf32>
// CHECK-LABEL: einsum_ellipsis

  // CHECK-LABEL @lowers_string_ops
  // CHECK-DAG: ml_program.global public @vars.Variable_1([]) : tensor<!tf_type.string>
  func.func @lowers_string_ops(%arg0: tensor<128xi32>, %arg1: tensor<128xi32>, %arg2: tensor<128x1xi32>, %arg3: tensor<128x90xi32>, %arg4: tensor<128x90xi32>, %arg5: tensor<128x90xi32>, %arg6: tensor<128x90x64xf32>, %arg7: tensor<128x90x64xf32>) -> tensor<!tf_type.string> {

  // CHECK:     [[ENTRY_SELECT:%.*]] = mhlo.select [[COMPARE_GT]], [[FEATURES]], [[ELSE_SELECT]]
  %0 = "tf.Softplus"(%arg0) : (tensor<8x16xf32>) -> tensor<8x16xf32>

  // CHECK:     return [[ENTRY_SELECT]] : tensor<8x16xf32>
  func.return %0 : tensor<8x16xf32>
}

// -----

// CHECK-LABEL: func @softplus_f64
// CHECK-SAME: ([[FEATURES:%.*]]: tensor<8x16xf64>)

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

    }) {device = "/device:CPU:0"} : () -> tensor<*xf32>
    // CHECK: "tf.Cast"(%{{.*}}) <{Truncate = false}> : (tensor<1x8x2xf32>) -> tensor<*xf32>
    // CHECK: (tensor<i32>, tensor<1x8x2xf32>) -> (tensor<1x8x1xf32>, tensor<1x8x1xf32>)
    %3:2 = "tf.Split"(%0, %1) {device = ""} : (tensor<i32>, tensor<*xf32>) -> (tensor<*xf32>, tensor<*xf32>)
    %4 = tensor.cast %1 : tensor<*xf32> to tensor<?x?x?xf32>

  %max_size = "tf.Const"() {value = dense<10> : tensor<i32>} : () -> tensor<i32>
  // CHECK-NOT: tf.EmptyTensorList
  // CHECK: tf.BroadcastTo
  // CHECK-SAME: tensor<10x16xf32>
  %tl0 = "tf.EmptyTensorList"(%element_shape, %max_size) : (tensor<1xi32>, tensor<i32>) -> tensor<!tf_type.variant<tensor<16xf32>>>
  // CHECK-NOT: tf.TensorListReserve
  // CHECK: tf.BroadcastTo

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

  func.return %0 : tensor<4x64x128xf32>
}

// CHECK-LABEL:   func @convert_gather_transpose(
// CHECK-SAME:                                   %[[VAL_0:.*]]: tensor<128x256xf32>,
// CHECK-SAME:                                   %[[VAL_1:.*]]: tensor<4x1xi32>) -> tensor<4x128xf32> {
// CHECK:           %[[VAL_2:.*]] = "tf.Const"{{.*}}value = dense<[1, 0]> : tensor<2xi64>