%0 = "tf.AddV2"(%arg1, %arg0) : (tensor<i32>, tensor<i32>) -> tensor<i32>
  func.return %0 : tensor<i32>
}

//===----------------------------------------------------------------------===//
// tf.XlaVariadicSort legalization
//===----------------------------------------------------------------------===//

// CHECK-LABEL: @xla_variadic_sort
// CHECK-SAME: %[[INPUT:.*]]: tensor<2x3x4xui8>

    BuildBodyWithCall(rewriter, loc, func, func_ty, &reduce_op.getBody());

    rewriter.replaceOp(op, reduce_op.getResults());

    return success();
  }
};

// Convert tf.XlaVariadicSort to mhlo.Sort
class ConvertXlaVariadicSortOp
    : public OpRewritePattern<TF::XlaVariadicSortOp> {
 public:
  using OpRewritePattern::OpRewritePattern;