namespace mlir {
namespace TF {
namespace {

// Returns true if the given op is TF/XLA communication op in the old bridge.
bool IsCommunicationOp(Operation* op) {
  return isa<TF::XlaHostComputeOp, TF::XlaSendToHostOp, TF::XlaRecvFromHostOp>(
      op);
}

// Returns true if the given op is one of ops supported to have communication
// subcomputation in the TF/XLA bridge.

  // MLIR. This will also require some trait to identify optional attributes in
  // MLIR.
  constexpr char kShapeInferenceGraph[] = "shape_inference_graph";
  if (mlir::isa<mlir::TF::XlaHostComputeOp>(inst) &&
      !inst->hasAttr(kShapeInferenceGraph) &&
      !attrs_to_ignore.contains(kShapeInferenceGraph)) {
    AttrValue value;
    value.mutable_func()->set_name("");

  supported_ops->insert(
      OperationName(mlir::TF::XlaCallModuleOp::getOperationName(), context));
  supported_ops->insert(
      OperationName(mlir::TF::XlaHostComputeOp::getOperationName(), context));
  supported_ops->insert(
      OperationName(mlir::TF::XlaReduceOp::getOperationName(), context));
  supported_ops->insert(

  // Infers the shape of function attached to XlaHostCompute.
  // Returns true if a return type was changed.
  bool InferShapeForFunctionAttachedToXlaHostCompute(XlaHostComputeOp op);

  // Infers the shape for MapDatasetOp and its associated function. Returns
  // whether either op or function type was changed.
  bool InferShapeForMapDataset(MapDatasetOp op, int64_t max_iterations);