if (element.getType() != slice_type) {
    update_slice = builder.create<TF::ReshapeOp>(
        loc, ArrayRef<Type>{slice_type},
        ArrayRef<Value>{element, GetR1Const(slice_shape, builder, loc)});
  }
  return builder
      .create<TF::XlaDynamicUpdateSliceOp>(
          loc, ArrayRef<Type>{buffer.getType()},
          ArrayRef<Value>{buffer, update_slice,
                          GetIndicesForElement(index, buffer, builder, loc)})

          << max_iterations_;
      return signalPassFailure();
    }
  }

 private:
  ArrayRef<ArrayRef<int64_t>> input_shapes_;
};
}  // namespace

std::unique_ptr<OperationPass<ModuleOp>> CreateTFShapeInferencePass(
    ArrayRef<ArrayRef<int64_t>> input_shapes) {
  return std::make_unique<ShapeInference>(input_shapes);
}

}  // namespace TF

// InferShapeForFunction.
FailureOr<bool> InferModuleShape(ModuleOp module, int64_t max_iterations = 10,
                                 ArrayRef<TypeID> ops_to_skip = {},
                                 ArrayRef<ArrayRef<int64_t>> input_shapes = {});

// Given a tensorflow NodeShape string, returns a vector of argument shapes
// that can be used with InferShapeForFunction.

// creating comm ops.
LogicalResult MoveToHostSingleCluster(
    ArrayRef<Operation*> clustered_ops, ArrayRef<Value> external_operands,
    ArrayRef<Value> external_outputs,
    llvm::SmallVector<IRMapping>& core_to_mapping,
    ArrayRef<Operation*> core_to_host_insertion_point,
    ArrayRef<Value> core_to_compilation_key,
    ArrayRef<Value> core_to_device_ordinal, StringAttr device_type_attr,

namespace tensorflow {

llvm::SmallVector<int64_t> ConvertTFShapeToMlir(llvm::ArrayRef<int64_t> shapes);

llvm::SmallVector<int64_t> ConvertMlirShapeToTF(llvm::ArrayRef<int64_t> shape);

static constexpr int64_t kTFDynamicSize = -1;
mlir::RankedTensorType GetTypeFromTFTensorShape(llvm::ArrayRef<int64_t> shape,
                                                mlir::Type elementType,

    mlir::ModuleOp module_op, mlir::MLIRContext* context,
    std::optional<Session*> session, bool run_tf_to_stablehlo,
    bool deserialize_xla_call_module,
    llvm::ArrayRef<llvm::ArrayRef<int64_t>> input_arg_shapes = {});

// Overload of `PreprocessAndFreezeGraph` that uses the default MLIR dump file
// prefix.
inline absl::Status PreprocessAndFreezeGraph(mlir::ModuleOp module_op,

  explicit MarkFunctionsNoinlinePass()
      : MarkFunctionsNoinlinePass(
            /*noinline_functions=*/ArrayRef<std::string>{}) {}

  // `noinline_functions` is a list of function names to be marked noinline.
  explicit MarkFunctionsNoinlinePass(
      const ArrayRef<std::string> noinline_functions)
      : noinline_functions_(CreateNoinlineFunctionsOption(noinline_functions)) {
  }

  // The following are metadata for inputs.
  llvm::ArrayRef<llvm::StringRef> input_names;
  llvm::ArrayRef<
      std::pair<tensorflow::DataType, tensorflow::PartialTensorShape>>
      input_specs;
  llvm::ArrayRef<llvm::StringRef> input_devices;

  // The following are metadata for outputs.
  llvm::ArrayRef<llvm::StringRef> output_names;
  llvm::ArrayRef<
      std::pair<tensorflow::DataType, tensorflow::PartialTensorShape>>

  template <typename Op>
  LogicalResult RewriteEqOp(Operation* op, PatternRewriter& rewriter) const;

  LogicalResult RewriteOp(
      Operation* op, PatternRewriter& rewriter,
      const std::function<bool(ArrayRef<int64_t>, ArrayRef<int64_t>,
                               SmallVectorImpl<int64_t>&)>&
          get_broadcasted_shape) const;

  LogicalResult RewriteBatchMatMulV2Op(Operation* op,

namespace mlir {
namespace TFL {

using llvm::ArrayRef;
using mlir::Operation;
using mlir::ShapedType;
using mlir::Value;

// Returns true if all tensor value in `values` has static shape and same shape.
inline bool OpHasSameStaticShapes(Operation* op) {
  auto values = op->getOperands();
  int operand_num = 0;
  ArrayRef<int64_t> shape;
  for (Value value : values) {