// 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,

// Rewrites computations generated by the xla.compile() Python code into
// XlaLaunch nodes.
//
// xla.compile() does two main things:
// a) marks operators that make up an XLA computation with the attribute
//    _xla_compile_id=XYZ, where XYZ is a unique key.
// b) adds XlaClusterOutput nodes to represent outputs of the computation.
//    These nodes are not marked with the _xla_compile_id attribute.

  func.return %0#0 : tensor<*xi32>
}

    $mlir_module is a serialized MLIR module with a `main` function that contains target computation.
    $metadata is a serialized TPUCompileMetadataProto describing the shapes and types of the inputs to the computation, as well as a mapping onto the TPU pod topology.

output.
### `-tf-tpu-host-computation-expansion`

_Expands host computation before and after TPU computation._

This pass expands outside compilation attributes to Identity/Cast ops
at the head of TPU computation if it's only used by outside compiled ops.
### `-tf-tpu-identity-pruning`

_Removes Identity/IdentityN ops from the TPU computation_

### `-tf-tpu-merge-variables-with-execute`

    %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>)

// RUN: tf-opt %s -split-input-file -tf-tpu-host-computation-expansion | FileCheck %s

// Tests expansion of a outside compiled ops at head/tail of TPU computation.

// CHECK-LABEL: func @identity_at_head_expanded
func.func @identity_at_head_expanded(%arg0: tensor<?xi32>) {
  // CHECK: "tf_device.cluster"
  // CHECK-NEXT: "tf.Identity"
  // CHECK-SAME: _xla_outside_compilation = ""
  "tf_device.cluster"() ({

       * to a valid final state ({@link #COMPLETED}, {@link #CANCELLED}, or {@link #INTERRUPTED}).
       *
       * @param v the value to set as the result of the computation.
       * @param t the exception to set as the result of the computation.
       * @param finalState the state to transition to.
       */
      private boolean complete(@Nullable V v, @Nullable Throwable t, int finalState) {

  // to execute an XLA Computation. Compilation results are cached. Compilation
  // is skipped if there is a cache hit. `function` is the name of a Tensorflow
  // function to compile. `args` is a description of the arguments to the
  // computation.
  //
  // `compile_mode` controls the behavior of the compilation cache on a cache

    cache.invalidateAll();
    // let the computation proceed
    computingLatch.countDown();
    // don't check cache.size() until we know the get("b") call is complete
    computationComplete.await();

    // At this point, the listener should be holding the seed value (a -> a), and the map should
    // contain the computed value (b -> b), since the clear() happened before the computation
    // completed.