// CHECK:       return %[[RES]]

func.func @torch_index_select(%arg0: tensor<2x1xf32>, %arg1: tensor<2xi32>) -> tensor<2x1xf32> {
  %0 = "mhlo.torch_index_select"(%arg0, %arg1) {
    batch_dims = 0 : i64, dim = 0 : i64
  } : (tensor<2x1xf32>, tensor<2xi32>) -> tensor<2x1xf32>
  func.return %0 : tensor<2x1xf32>
}

// CHECK-LABEL:   func @lowered_cumsum(

  %2057 = "tfl.transpose"(%arg0, %cst_10) : (tensor<1x16x256xf32>, tensor<3xi32>) -> tensor<1x256x16xf32>
  %2058 = "tfl.reshape"(%2057, %cst_3) : (tensor<1x256x16xf32>, tensor<2xi32>) -> tensor<256x16xf32>
  %2059 = "tfl.transpose"(%2058, %cst_6) : (tensor<256x16xf32>, tensor<2xi32>) -> tensor<16x256xf32>
  return %2059: tensor<16x256xf32>
  // CHECK-DAG: %cst = arith.constant dense<[16, 256]> : tensor<2xi32>

func.func @alltoall_basic(%input: tensor<1x2xf32>) -> tensor<2x1xf32> {
  %group_assignment = "tf.Const" () {
    value = dense<[[0, 1]]> : tensor<1x2xi32>
  } : () -> tensor<1x2xi32>
  %result = "tf.AllToAll"(%input, %group_assignment) {T = f32, concat_dimension = 0 : i64, split_count = 2 : i64, split_dimension = 1 : i64} :  (tensor<1x2xf32>, tensor<1x2xi32>)  -> tensor<2x1xf32>
  // CHECK: mhlo.all_to_all