if (n->type_string() == "_XlaRecvAtHost" ||
        n->type_string() == "_XlaSendFromHost") {
      n->ClearAttr("device_ordinal");
      n->AddAttr("device_ordinal", device_ordinal_value);
    } else if (n->IsIfNode()) {
      for (const string& attr_name :
           std::vector<string>{"then_branch", "else_branch"}) {
        NameAttrList branch_func;
        TF_RETURN_IF_ERROR(GetNodeAttr(n->attrs(), attr_name, &branch_func));

  // Map Case/If/While and StatelessCase/If/While op in TensorFlow to the common
  // Case/If/While op in MLIR and add the differentiating attribute.
  if (node.IsCaseNode()) composite_control_flow_op("Case");
  if (node.IsIfNode()) composite_control_flow_op("If");
  if (node.IsWhileNode()) {
    composite_control_flow_op("While");
    auto* output_shapes = node.attrs().Find("output_shapes");