// CHECK: return %[[v5]] : tensor<?x?x8xf32>
}

func.func @einsum_no_reshape(%arg0 : tensor<?x?x8x128xf32>, %arg1 : tensor<1x?x8x128xf32>) -> tensor<?x?x8x?xf32> {
  %0 = "tf.Einsum"(%arg0, %arg1) {device = "", equation = "BQNH,BTNH->BQNT"} : (tensor<?x?x8x128xf32>, tensor<1x?x8x128xf32>) -> tensor<?x?x8x?xf32>
  func.return %0 : tensor<?x?x8x?xf32>
// CHECK-LABEL: einsum_no_reshape

   "tfl.yield"(%4) : (tensor<128x128xf32>) -> ()
  }) {_tfl_quant_trait = "fully_quantizable", device = ""} : (tensor<128x128xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xi32>) -> tensor<128x128xf32>

   "tfl.yield"(%4) : (tensor<128x128xf32>) -> ()
  }) {_tfl_quant_trait = "fully_quantizable", device = ""} : (tensor<128x128xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xi32>) -> tensor<128x128xf32>

}

// -----

module {
func.func private @func_20_GPU_FLOAT(%arg0: tensor<128x128xf32>, %arg1: tensor<3xi32>) -> tensor<1x128x128xf32> attributes {tac.device = "GPU", tac.inference_type = "FLOAT", tac.interface_name = "func_20"} {
  %0 = "tfl.reshape"(%arg0, %arg1) {tac.device = "GPU", tac.inference_type = "FLOAT"} : (tensor<128x128xf32>, tensor<3xi32>) -> tensor<1x128x128xf32>
  func.return %0 : tensor<1x128x128xf32>
}

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