namespace tensorflow {
namespace tf2xla {
namespace internal {
namespace {

using OpShardingVariant = std::variant<mlir::Operation*, llvm::StringRef>;
using OpShardingVector = llvm::SmallVector<OpShardingVariant, 8>;
using OptionalOpShardingVector =
    llvm::SmallVector<std::optional<OpShardingVariant>, 8>;
using llvm::StringRef;
using mlir::Block;
using mlir::BlockArgument;
using mlir::BoolAttr;