}

// Returns a set of ops that are outside compiled and can be extracted to before
// the TPU computation. These ops are either connected to the inputs of the TPU
// computation or other ops that can be extracted, and have no operands from
// other ops in the TPU computation that cannot be extracted.
llvm::SmallVector<Operation*, 4> FindOutsideCompiledOpsAtHead(

// then those passes include all the Legalization to XLA HLO which is returned
// in the compilation_result.
absl::StatusOr<std::string> CompileFromMlirToXlaHlo(
    bool lower_to_xla_hlo, const tpu::MlirToHloArgs& computation,
    const tpu::TPUCompileMetadataProto& metadata, llvm::StringRef device_type,
    const XlaShapeLayoutHelpers::ShapeDeterminationFns& shape_determination_fns,

    %add_result = "tf.Add"(%computation, %computation) : (tensor<i32>, tensor<i32>) -> tensor<i32>
    // expected-error@+1 {{Input of TPUPartitionedOutputV2 must be in tpu computation.}}
    %partitioned_output:2 = "tf.TPUPartitionedOutputV2"(%add_result) {N = 2 : i64, partition_dims = []} : (tensor<i32>) -> (tensor<i32>, tensor<i32>)

namespace tf2xla {
namespace v2 {

// Legalizes the given mlir::Module into XLA HLO. If successful, returns the
// compiled XLA HLO. V1 of the tf2xla uses MLIR whereas V0 does not use MLIR.
//
// Inputs:
//  computation - The MLIR module op. It currently takes in
//  tpu::FunctionToHloArgs but this is deprecated. arg_shapes - The shapes of
//  the arguments in module_op. device_type - The device type to compile for.

CreateTPUClusterFormationPass(bool strict_clusters = false);

// Creates a pass that extracts outside compilation (Host ops inside device
// cluster) at head/tail of Device cluster to run before/after XLA computation.
std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
CreateExtractHeadTailOutsideCompilationPass();

// Creates a pass that extract outside compilation (Host ops inside cevice

    std::vector<std::vector<xla::Shape>> per_core_arg_shapes;

    tpu::TPUCompileMetadataProto metadata_proto;
    std::vector<TensorShape> arg_shapes;
    if (computation.index() == 0) {
      TF_RETURN_IF_ERROR(tensorflow::tf2xla::internal::ConfigureMetadata(
          std::get<0>(computation).mlir_module, arg_shapes, metadata_proto));
    }

    XlaShapeLayoutHelpers::ShapeDeterminationFns shape_determination_fns;

// it's double-declared (it would cost extra in the common case to verify
// this). But the MethodSet computation will not find it due to the name
// collision caused by the double-declaration, leading to an internal
// inconsistency while we are verifying one computation against the other.
// var _ = T1c{}.Pointer

// T2's method declared before the type.
func (*T2) f /* ERROR "field and method" */ () {}

{{DECLS_FROM_OBJ_FILE}}

{{NS_START}}
// {{CLASS}} represents a computation previously specified in a
// TensorFlow graph, now compiled into executable code. This extends the generic
// XlaCompiledCpuFunction class with statically type-safe arg and result
// methods. Usage example:
//
//   {{CLASS}} computation;
//   // ...set args using computation.argN methods
//   CHECK(computation.Run());

// use_tuple_args: when this is true, always create a tuple argument for the
//   entry computation.
// enable_op_fallback: when this is true, prefer tf2xla fallback kernels over
// MLIR
//   native kernels for legalization to HLO.
// return_tuple: when this is true, always create a tuple result for the
//   entry computation.
// shape_determination_fns: Contains layout preference fn and shape

    }
  }
}

// Since we have the outputs from host and device computation after moving
// outside compiled ops, we can create the actual parallel_execute regions.
// Still, one region is for the host computation for outside compilation and
// the other one is for the original Device cluster computation.
mlir::tf_device::ParallelExecuteOp CreateFinalParallelExecuteOp(