including:

* TensorFlow IR, which represents all things possible in TensorFlow graphs.
* XLA HLO IR, which is designed to take advantage of XLA’s compilation
  abilities (with output to, among other things, TPUs).
* An experimental affine dialect, which focuses on
  [polyhedral representations](https://en.wikipedia.org/wiki/Polytope_model)
  and optimizations.
* LLVM IR, which has a 1:1 mapping between it and LLVM’s own representation,

                          WeightOnlyPtqComponent::kName, *function_aliases,
                          *ctx, *module));

  // Remove the `tpu` tag for exporting because the output quantized model is
  // essentially a CPU model.
  tags.erase("tpu");

  py_function_library.SaveExportedModel(
      dst_saved_model_path, post_calibrated_exported_model,
      src_saved_model_path, tags, signature_def_map);

  return builder->create<mlir::TF::ConcatOp>(
      location, output_type, concat_dimension_op.getOutput(), inputs);
}

// For tile sharded inputs to TPU computation, inject split op between the
// input values and TPU computation so that tiled input values are passed in
// as inputs to TPU computations. If more than one dimension is sharded, then
// a tree of connected split ops are added before tf_device.parallel_execute op.

# Script for mlir-reduce to test input interestingness.
tf-opt $1 -tf-mark-ops-for-outside-compilation &> | grep "'tf.UnsupportedOp' op isn't compilable for TPU device"

if [ $? -eq 0 ]; then
  exit 1
else
  exit 0

                          PostCalibrationComponent::kName, *function_aliases,
                          *ctx, *module));

  // Remove the `tpu` tag for exporting because the output quantized model is
  // essentially a CPU model.
  tags.erase("tpu");

  py_function_library.SaveExportedModel(
      dst_saved_model_path, post_calibrated_exported_model,
      src_saved_model_path, tags, signature_def_map);

    "tf_device.cluster_func"(%arg0) <{func = @empty_func}> {_dynamic_arg_index = [], _replication_info = "cluster", _xla_compile_device_type = "TPU", allow_soft_placement = false, computation_shape = [], device = "", device_assignment = [0, 0, 0, 0, 0, 0, 0, 1], host_compute_core = [], input_sharding_configuration = ["{devices=[2,1]0,1}"], num_cores_per_replica = 2 : i64, output_sharding_configuration...

// RUN: tf-opt %s -tf-tpu-colocate-composite-resource-ops | FileCheck %s

// Tests ReadVariable op using composite device resource is wrapped inside
// tf_device.Cluster.

// CHECK-LABEL: func @testReadVariableOpColocated
// CHECK-SAME: (%[[ARG0:.*]]: tensor<*x!tf_type.resource<tensor<4xf32>>>)
func.func @testReadVariableOpColocated(%arg0: tensor<*x!tf_type.resource<tensor<4xf32>>>) {
  // CHECK:      tf_device.replicate

// RUN: tf-opt %s -allow-unregistered-dialect --tf-move-tpu-compile-to-front --split-input-file | FileCheck %s

module {

// CHECK-LABEL: does_basic_reordering
func.func @does_basic_reordering() -> () {
   // CHECK: _TPUCompileMlir
   // CHECK-SAME: X
   // CHECK: _TPUCompileMlir
   // CHECK-SAME: Y
   // CHECK: OpA
   // CHECK: OpB
   // CHECK: OpC
   "tf.OpA"() : () -> ()

#include "mlir/IR/OwningOpRef.h"  // from @llvm-project

namespace mlir {

// Set layouts attribute of tf.InfeedDequeueTuple ops.
bool SetTPUInfeedLayout(OwningOpRef<ModuleOp>& mlir_module);

// Try to determine the right TPU infeed layout.
FailureOr<Attribute> GetTPUInfeedLayout(ArrayRef<Type> types,
                                        OpBuilder& rewriter);

}  // namespace mlir

#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/attribute_utils.h"
#include "xla/mlir_hlo/mhlo/IR/hlo_ops.h"

#define DEBUG_TYPE "tf-tpu-annotate-dynamic-shape-inputs"

namespace mlir {
namespace TFTPU {

namespace {

#define GEN_PASS_DEF_TPUANNOTATEDYNAMICSHAPEINPUTSPASS
#include "tensorflow/compiler/mlir/tensorflow/transforms/tf_passes.h.inc"