}
      subgraph.available_choices.emplace(*inference_device_type, impl);
    }

    // Insert in the subgraphs.
    subgraphs_.try_emplace(call_op, subgraph);

    // If it's an output subgraph, we will add to the output_subgraphs.
    if (returned_call_op_set.find(call_op) != returned_call_op_set.end()) {
      output_subgraphs->push_back(call_op);
    }
  });
}

    // Insert all ops including nested ops for checking outputs/side effects.
    cluster_op.walk(
        [&](Operation* op) { tail_outside_compiled_ops_set.insert(op); });

    // Only add top level ops to output vector.
    tail_outside_compiled_ops->push_back(&cluster_op);
  }

  *cluster_results = cluster_results_set.takeVector();
}

    return success();
  }
};

void SessionInitializerOp::getCanonicalizationPatterns(
    RewritePatternSet &results, MLIRContext *context) {
  results.add<OptimizeSessionInitializerPattern>(context);
}

SmallVector<StringRef, 2> GetSessionInitializerExportedName(ModuleOp op) {
  auto session_initializer_op = GetSessionInitializerOp(op);

        TypeID::get<TF::XlaSpmdFullToShardShapeOp>(),
        TypeID::get<TF::XlaSpmdShardToFullShapeOp>(),
        TypeID::get<TF::XlaSvdOp>(),
    };

    // Add the ops from the TPUEmbeddingOpsRegistry.
    for (auto op_type_id :
         TF::TPUEmbeddingOpsRegistry::Global().GetOpsTypeIds()) {
      ops_set->insert(op_type_id);
    }
    return ops_set;
  }();

               std::back_inserter(result), [](const BufferInfo& buffer_info) {
                 return buffer_info.is_temp_buffer();
               });
  return result;
}

// Add (from,to) rewrite pairs based on the given shape.  These rewrite pairs
// are used to generate methods for args and results.
Status AddRewritesForShape(int i, const xla::Shape& shape,

// If GCS_APPEND_MODE=compose then instead the new data is uploaded to a
// temporary object and composed with the original object. This is disabled by
// default as the multiple API calls required add a risk of stranding temporary
// objects.
constexpr char kComposeAppend[] = "compose";

// We can cast `google::cloud::StatusCode` to `TF_Code` because they have the
// same integer values. See