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,

    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,

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

static SmallVector<int64_t> InvertPermutationToVector(ArrayRef<int64_t> vec,
                                                      ArrayRef<int64_t> perm) {
  return applyPermutation(vec, invertPermutationVector(perm));
}

static RankedTensorType GetPermutedTensorTypeHelper(RankedTensorType type,
                                                    ArrayRef<int64_t> perm,

    }
    VarHandleOp var_handle_op = builder.create<VarHandleOp>(
        variable_v2_op.getLoc(),
        ArrayRef<Type>{RankedTensorType::get(
            {}, TF::ResourceType::get(ArrayRef<TensorType>{tensor_type},
                                      builder.getContext()))},
        ArrayRef<Value>{},
        ArrayRef<NamedAttribute>{
            builder.getNamedAttr("device", device_attr),
            builder.getNamedAttr("container",

  return ArrayAttr::get(attr.getContext(), shuffled);
}

// Shuffle ranked tensor dimensions according to the permutation.
Type ShuffleRankedTensorType(Type type, ArrayRef<int64_t> permutation) {
  if (auto ranked_type = mlir::dyn_cast<RankedTensorType>(type)) {
    ArrayRef<int64_t> shape = ranked_type.getShape();
    assert(permutation.size() == shape.size());

    SmallVector<int64_t, 4> new_shape(permutation.size());