transposed_shape[i] = input_type.getDimSize(permutation_array[i]);
  }
  auto transposed_type =
      RankedTensorType::get(transposed_shape, input_type.getElementType());
  DenseIntElementsAttr permutation = DenseIntElementsAttr::get(
      RankedTensorType::get(permutation_array.size(), rewriter.getI64Type()),
      permutation_array);
  return {permutation, transposed_type};
}

  if (!mlir::isa<RankedTensorType>(input_lhs.getType())) {
    // LHS must be a ranked tensor type
    return failure();
  }
  if (!mlir::isa<RankedTensorType>(input_rhs.getType())) {
    // RHS must be a ranked tensor type
    return failure();
  }

  auto lhs_type = mlir::cast<RankedTensorType>(input_lhs.getType());
  auto rhs_type = mlir::cast<RankedTensorType>(input_rhs.getType());

      lhs_dot_dimensions_info.contracting_dimensions().SizesArray());
  auto lhs_transposed = rewriter.create<mhlo::TransposeOp>(
      loc,
      RankedTensorType::get(lhs_transposed_shape, lhs_type.getElementType()),
      lhs,
      DenseIntElementsAttr::get(
          RankedTensorType::get({lhs_rank}, rewriter.getI64Type()),
          lhs_permutation));

  // Transposes rhs shape to be in the order of {batch_dimensions, contracting

    auto weight = op->getOperand(1);
    auto weight_type =
        mlir::dyn_cast_or_null<mlir::RankedTensorType>(weight.getType());
    if (weight_type == nullptr || !weight_type.hasStaticShape()) return false;

    auto output = op->getResult(0);
    auto output_type =
        mlir::dyn_cast_or_null<mlir::RankedTensorType>(output.getType());
    if (output_type == nullptr || !output_type.hasStaticShape()) return false;

    if (padding == "EXPLICIT" && num_spatial_dims == 3) {
      auto lhs_type =
          mlir::dyn_cast<RankedTensorType>(conv_op.getLhs().getType());
      RankedTensorType padding_attr_type = mlir::RankedTensorType::get(
          {lhs_type.getRank(), 2}, rewriter.getIntegerType(64));
      auto padding_const = rewriter.create<TF::ConstOp>(
          conv_op->getLoc(),

void HandleFirstConvolution(TF::Conv2DOp conv2d, int64_t block_size) {
  // Check if input and filter type are RankedTensorType.
  auto input_tensor_type =
      mlir::dyn_cast<RankedTensorType>(conv2d.getInput().getType());
  auto filter_tensor_type =
      mlir::dyn_cast<RankedTensorType>(conv2d.getFilter().getType());
  if (!input_tensor_type || !filter_tensor_type) return;

                                    SmallVectorImpl<int64_t>* out_shape,
                                    PatternRewriter* rewriter) {
  RankedTensorType lhs_type = mlir::cast<RankedTensorType>(lhs->getType());
  RankedTensorType rhs_type = mlir::cast<RankedTensorType>(rhs->getType());

  int32_t num_lhs_reshape_segids = 0;
  int32_t num_rhs_reshape_segids = 0;
  std::vector<int32_t> lhs_reshape_segids;
  int lhs_rank =

// type static shaped RankedTensorType. `unique_index` is used to get the unique
// value for the attribute.
mlir::ElementsAttr GetSplat(mlir::RankedTensorType type, int unique_index,
                            mlir::Builder builder);

absl::StatusOr<mlir::ElementsAttr> ConvertIntBuffer(
    mlir::RankedTensorType shaped_type, const std::vector<uint8_t>& buffer,
    bool truncate = false);

namespace mlir {

class Builder;

namespace TF {

// Returns the RankedTensorType for the given operand. TensorFlow constant ops
// may have non-static shape because the shape is not propagated during constant
// folding. If the defining op for the given operand is a constant op, this
// routine uses the constant op's attribute to get the actual shape.
RankedTensorType GetRankedTensorTypeForOperand(Value operand);

  "DenseElementsAttr::get(RankedTensorType::get({1}, $_builder.getIntegerType(32)), "
  "static_cast<int32_t>($0.getType().cast<RankedTensorType>().getDimSize(  "
  "  $0.getType().cast<RankedTensorType>().getRank() - 1)))">;

// Check whether the tensor is ranked and whether its last dim is static.
def IsRankedShapeLastDimStatic : Constraint<And<[
  CPred<"$0.getType().isa<RankedTensorType>()">,