return shaped_type.clone(new_shape);
}

// Returns a ShapedType for a permutation and the shape of input after
// applying the permutation to the given shape through a transpose.
inline ShapedType GetTransposedType(Value input,
                                    llvm::ArrayRef<int64_t> permutation_array) {
  auto input_type = input.getType().cast<ShapedType>();
  if (permutation_array.size() != input_type.getRank()) {

// indices in `val2`.
bool HasEqualElementSize(Value val1, Value val2, ArrayRef<int> val1_indices,
                         ArrayRef<int> val2_indices) {
  ShapedType val1_shape = mlir::cast<ShapedType>(val1.getType());
  ShapedType val2_shape = mlir::cast<ShapedType>(val2.getType());
  if (!val1_shape.hasRank() || !val2_shape.hasRank()) return false;

  int val1_result = 1;
  int val2_result = 1;

// logics.
std::optional<TF::PartitionedCallOp> RegisterOperationsInFuncOp(
    StringRef func_name, PatternRewriter& rewriter, QuantizedType quant_type,
    Value input_val, ShapedType result_type,
    std::function<Operation*(PatternRewriter&, Operation*, Value, ShapedType,
                             QuantizedType)>
        quantization_operations_func) {
  Operation* input_op = input_val.getDefiningOp();

    const float ratio_threshold) {
  ElementsAttr attr;
  ShapedType type;
  InspectResult result = {};
  if (auto cst = dyn_cast<ConstOp>(inst)) {
    attr = cst.getValue();
    type = mlir::cast<ShapedType>(cst.getType());
  } else if (auto cst = dyn_cast<QConstOp>(inst)) {
    attr = cst.getValue();
    type = mlir::cast<ShapedType>(cst.getType());
  } else {
    result.can_compress = false;

    if (!mlir::cast<ShapedType>(lhs.getType()).hasStaticShape() ||
        !mlir::cast<ShapedType>(rhs.getType()).hasStaticShape() ||
        !mlir::cast<ShapedType>(cond.getType()).hasStaticShape()) {
      return rewriteOpWithDynamicInput(op, rewriter);
    }

    auto lhs_shape = mlir::cast<ShapedType>(lhs.getType()).getShape();
    auto rhs_shape = mlir::cast<ShapedType>(rhs.getType()).getShape();

      if (operands.size() != 1 || updates.size() != 1) return failure();

      ShapedType operand_type = mlir::cast<ShapedType>(operands[0].getType());
      ShapedType indices_type = mlir::cast<ShapedType>(indices.getType());
      ShapedType updates_type = mlir::cast<ShapedType>(updates[0].getType());

      Value new_updates = updates[0];

      // Can only convert with static shaped scatter.

Value ConvertDot(PatternRewriter& rewriter, Value lhs, Value rhs,
                 mhlo::DotDimensionNumbersAttr dot_dimension_numbers,
                 ShapedType result_type, mlir::Location loc) {
  auto lhs_type = mlir::cast<ShapedType>(lhs.getType());
  auto rhs_type = mlir::cast<ShapedType>(rhs.getType());
  const int lhs_rank = lhs_type.getRank();
  const int rhs_rank = rhs_type.getRank();
  ImplicitLocOpBuilder builder(loc, rewriter);

void BroadcastBatchDimensionsForBatchMatMul(OpBuilder &builder, Location loc,
                                            Value &input, Value &weight) {
  ShapedType input_type = mlir::cast<ShapedType>(input.getType());
  ShapedType weight_type = mlir::cast<ShapedType>(weight.getType());
  const int32_t input_rank = input_type.getRank();
  const int32_t weight_rank = weight_type.getRank();

    return builder.create<TF::ReshapeOp>(
        loc, RankedTensorType::get(shape, builder.getI32Type()), value,
        CreateConstValue<int64_t>(builder, loc, {rank}, shape));
  };

  ShapedType filter_shape = mlir::cast<ShapedType>(filter.getType());
  Value input_shape_value = builder.create<TF::ShapeOp>(
      loc, RankedTensorType::get({num_dims}, builder.getI32Type()), input);

  dot_dimension_numbers.ParseFromString(dot_dimension_numbers_str.str());
  SmallVector<Value> input_arguments = {lhs, rhs};
  const int lhs_rank = mlir::cast<ShapedType>(lhs.getType()).getShape().size();
  const int rhs_rank = mlir::cast<ShapedType>(rhs.getType()).getShape().size();

  const std::string einsum_equation =
      CreateEinsumEquation(dot_dimension_numbers, lhs_rank, rhs_rank);