auto shape = ranked_ty.getShape();
  int rank = shape.size();

  SmallVector<APInt, 4> dimensions;
  dimensions.reserve(rank);
  for (int i = 0; i < rank; ++i)
    dimensions.push_back(APInt(out_width, shape[i]));

  auto result_type = tensorflow::GetTypeFromTFTensorShape(
      {rank}, IntegerType::get(input_ty.getContext(), out_width));

  }

  const int64_t rank = input_type.getRank();
  if (rank <= 0) {
    return emitOptionalError(loc, "input should be of rank larger than 0");
  }

  int64_t axis_value = op.getAxisAttr().getInt();
  if (axis_value < 0) {
    axis_value += rank;
  }
  if (axis_value < 0 || axis_value >= rank) {
    return emitOptionalError(

  if (!input_ty) return success();

  int rank = input_ty.getRank();
  if (rank != 1 && rank != 2)
    return op.emitOpError("requires input of rank 1 or 2");

  if (rank == 1) {
    int64_t dim0 = input_ty.getDimSize(0);
    if (dim0 != ShapedType::kDynamic && dim0 != 4 && dim0 != 2)
      return op.emitOpError("requires 1D input of size 4 or size 2");
  }

  if (rank == 2) {

      input = InsertDilateOp(op, rewriter);
    }

    TensorType operand_type = input.getType().cast<TensorType>();
    const int64_t rank = operand_type.getRank();
    // Shape of padding should be [rank, 2].
    SmallVector<int64_t> shape{rank, 2};
    TensorType padding_type =
        operand_type.cloneWith(shape, rewriter.getI32Type());

    ArrayRef<int64_t> padding_low = op.getEdgePaddingLow();

  func.return %0 : tensor<2x2xf32>
}

// -----

// Test tf.Transpose with invalid rank of y
func.func @testTranspose(tensor<2x3xf32>) -> tensor<3x2x1xf32> {
^bb0(%arg0: tensor<2x3xf32>):
  %cst = arith.constant dense<[1, 0]> : tensor<2xi32>
  // expected-error @+1 {{x should be of the same rank with y, got x of rank 2, and y of rank 3}}

// location.
// Returns true if the n-th operand has unknown rank or has rank m.
class TFL_OperandHasRank<int n, int m> :
  PredOpTrait<"operand " # n # " is " # m # "-D",
    Or<[TFL_OperandIsUnrankedPred<n>,
      CPred<"$_op.getOperand(" # n #
      ").getType().cast<ShapedType>().getRank() == " # m>]>>;

// Returns true if the n-th operand is ranked and has rank dim.
class TFL_OperandHasKnownRank<int n, int dim> : And<[

  }

  tensorflow::shape_inference::ShapeHandle shape = ic->output(output.index);

  int rank = -1;
  if (ic->RankKnown(shape)) {
    rank = ic->Rank(shape);
  }

  if (num_dims != rank) {
    status->status = InvalidArgument("Expected rank is ", num_dims,
                                     " but actual rank is ", rank);
    return;
  }

  if (num_dims == 0) {
    // Output shape is a scalar.

  bool TransposeFirstTwoDimToLast(DenseIntElementsAttr perm) const {
    int rank = perm.getNumElements();
    if (rank < 3) return false;
    for (int i = 0; i < rank - 2; i++) {
      if (perm.getValues<APInt>()[i] != i + 2) {
        return false;
      }
    }
    return perm.getValues<APInt>()[rank - 2] == 0 &&
           perm.getValues<APInt>()[rank - 1] == 1;
  }
};

  func.return %0 : tensor<2x2xi32>
}

// -----

func.func @pack(%arg0: tensor<2xi32>, %arg1: tensor<2xi32>) -> tensor<2x2xi32> {
  // expected-error @+1 {{op attribute 'axis' should be in range [-rank - 1, rank + 1), got rank = 1, and axis = 3}}
  %0 = "tfl.pack"(%arg0, %arg1) {axis = 3 : i32, values_count = 2 : i32} : (tensor<2xi32>, tensor<2xi32>) -> tensor<2x2xi32>
  func.return %0 : tensor<2x2xi32>
}

// -----

    auto sort_dim = op.getDimension();
    auto k = indices_ty.getDimSize(sort_dim);
    auto rank = keys_ty.getRank();
    if (sort_dim != rank - 1 || k < 1)
      return rewriter.notifyMatchFailure(
          op, "only match for sort dim = rank - 1 and DimSize >= 1");

    // In the following, we'll check indices is obtained by a iota.
    auto sort_dim_attr = DenseIntElementsAttr::get(