}) : () -> tensor<f32>
      return %0 : tensor<f32>
    }
    ```
  }];

  let constructor = "tensorflow::tf2xla::internal::CreateExtractOutsideCompilationPass()";
}

def XlaClusterFormationPass : Pass<"tf-xla-cluster-formation", "ModuleOp"> {
  let summary = "Encapsulate partitioned calls within a Cluster op";
  let description = [{
    This pass clusters `tf.PartitionedCall` and `tf.StatefulPartitionedCall`