void runOnFunction(
      func::FuncOp func,
      const TF::ResourceAliasAnalysis::Info& resource_alias_analysis);

  StringRef getArgument() const final { return "tf-tpu-dynamic-layout-pass"; }

  StringRef getDescription() const final {
    return "Inserts TPU layout ops to determine layout at run time.";
  }
};

// Checks if the input producer op is supported in this transform. Right now, we

    return absl::OkStatus();
  }

  // TFRT-TPU is used if device type is `DEVICE_TPU` and platform_info does not
  // have `xla_device_metadata`. This is used for TFRT-TPU when
  // BuildXlaDeviceCompiler() is called in GetCompilerIr(). Currently only
  // lowering to HLO is needed there and xla::LocalClient doesn't support
  // building the executable for TFRT-TPU and hence, is set to nullptr here.

  let summary = "Op that compiles a computation in MLIR into a TPU program, and loads and executes it on a TPU device.";

  let description = [{
For the internal use of the TPU compiler.

'static_shapes' are tensors specifying the maximum dimension sizes for the tensors specified in `dynamic_operands`.
'args' are inputs to the TPU computation.

  // CHECK: %[[IDENTIFY:.*]] = "tf.Identity"(%[[SUBGRAPH_0]]#1) {device = ""} : (tensor<1024x3xf32>) -> tensor<1024x3xf32>
  // CHECK: %[[SUBGRAPH_1:.*]] = "tf.XlaCallModule"() <{Sout = [#tf_type.shape<1024x3>], {{.*}} ["CPU", "TPU"], {{.*}}}> {_entry_function = @_stablehlo_main_1

                          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.

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

#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"

        "//tensorflow/compiler/tf2xla/kernels:xla_ops",
        "//tensorflow/core:framework",
        "//tensorflow/core:test_main",
        "//tensorflow/core/protobuf/tpu:compile_metadata_proto_cc",
        "//tensorflow/core/tpu/kernels/xla:host_compute_ops",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/strings:string_view",
        "@llvm-project//mlir:FuncDialect",
        "@llvm-project//mlir:IR",

// RUN: tf-tfrt-opt -split-input-file -tf-mlrt-rewrite-ifrt-load-variable %s | FileCheck %s

// Variable is used by both CPU and TPU
//
// CHECK-LABEL: func @serving_default(%arg0: tensor<1x3xf32>) -> tensor<1x1xf32>
// CHECK-NEXT:    [[HANDLE:%.*]] = "tf.VarHandleOp"()
// CHECK-NEXT:    [[ARRAYKEY:%.*]], [[FURTURE:%.*]] = "tf_mlrt.tf_ifrt_load_variable"([[HANDLE]])