$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.

// are specified in `variable_indices` argument. We snapshot tensors that back
// resource variables since concurrent updates may modify the shape, and it is
// important that the shapes used for compilation match the true shapes of the
// buffers.
//
// We snapshot the entire set of resource variables as one atomic operation.
// This models Read->* dependencies between resource variable operations.  See

 public:
  CheckpointReader(const string& filename, TF_Status* status);

  bool HasTensor(const string& name) const;
  const string DebugString() const;

  // Returns a map from variable names to their shapes.  Slices of a partitioned
  // tensor are combined into a single entry.
  const TensorSliceReader::VarToShapeMap& GetVariableToShapeMap() const;

  // Returns a map from variable names to their data types.  Slices of a

                                         /*shape=*/shape_tensor);
}

// Creates ReshapeOp with runtime calcuation of required shape to support
// dynamic shapes. The shape is calculated by Shape and UnsortedSegmentProd op.
// `reshape_segids` and `num_reshape_segids` for UnsortedSegmentProd is
// calculated in `reshapeForBatchMatmul`.
TF::ReshapeOp createReshapeOpForDynamic(Value value, ArrayRef<int64_t> shape,

      location, concat_dim_type, concat_dimension_attr);

  // Correctly set output shapes of concat op output if output shape is
  // statically known. Since the shape of TPUExecute op must be the same
  // across logical devices, we refer to the shape of 0th logical device
  // computation output.
  mlir::Type output_type;
  auto input_type = mlir::cast<mlir::TensorType>(inputs[0].getType());

    if (ShapedType::isDynamic(num_index_dims)) return failure();

    auto updates = op.getUpdates();

    // Broadcast scalar `updates` in into expected shape as following shape:
    // updates.shape == indices.shape[:-1] + tensor.shape[indices.shape[-1]:]
    if (updates_ty.getRank() == 0 &&
        (std::is_same<OpTy, TF::TensorScatterUpdateOp>::value ||
         std::is_same<OpTy, TF::TensorScatterAddOp>::value)) {

        /*batch_group_count=*/op.getBatchGroupCountAttr(),
        /*precision_config=*/op.getPrecisionConfigAttr());

    // Transpose the output of the `ConvolutionOp` back to the original op's
    // output shape so that users' shapes match.
    // [b, 0, 1, f] => [b, f, 0, 1]
    auto output_transpose_op = rewriter.create<mlir::stablehlo::TransposeOp>(
        new_convolution_op.getLoc(), /*resultType0=*/output_tensor_type,

  // CHECK-NEXT: %[[SHAPE0:.*]] = shape.shape_of %arg0 : tensor<?xi1> -> tensor<1xindex>
  // CHECK-NEXT: %[[SHAPE1:.*]] = shape.shape_of %arg1 : tensor<?x?x8xi32> -> tensor<3xindex>
  // CHECK-NEXT: %[[SHAPE2:.*]] = shape.shape_of %arg2 : tensor<?x?x8xi32> -> tensor<3xindex>
  // CHECK-NEXT: %[[SHAPEEQ1:.*]] = shape.cstr_eq %[[SHAPE1]], %[[SHAPE2]] : tensor<3xindex>, tensor<3xindex>

namespace mlir {
namespace TF {

namespace {

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

// This transformation pass propagate shapes on the TensorFlow graph.
// It is a ModulePass in order to be able to change function types.
class ShapeInference
    : public impl::TensorFlowShapeInferencePassBase<ShapeInference> {
 public:
  ShapeInference() = default;

        return mlir::UnrankedTensorType::get(
            mlir::TF::ResourceType::get(builder.getContext()));
      }
    } else if (auto shape = node.attrs().Find("_output_shapes")) {
      if (shape->has_list() && shape->list().shape_size() == 1) {
        return ConvertToMlirTensorType(shape->list().shape().at(0), dtype,
                                       &builder);
      }
    }
  }