(getElementTypeOrSelf(op.getOutput().getType()))))
      return failure();

    ElementsAttr input_tensor = qconst_op.getValue();

    assert(perm_tensor.getType().getRank() == 1);
    const int num_dimensions = input_tensor.getShapedType().getRank();
    assert(perm_tensor.getType().getNumElements() == num_dimensions);

    ArrayRef<int64_t> input_shape = input_tensor.getShapedType().getShape();

    if (std::isalpha(label)) {
      new_lhs.push_back(label);
    } else {
      // Encounter ellipsis: generate unnamed labels then insert to the new
      // labels.
      new_labels = GenerateLabels(lhs_ty.getRank() - lhs_named_label_count,
                                  available_labels);
      new_lhs.append(new_labels);
      i += 2;
    }
  }

  std::string new_rhs, new_rhs_labels;

    // and {num_channels} (rank 1) for per-channel quantized one.
    auto scale_type = filter_scale_attr.getType().dyn_cast<RankedTensorType>();
    if (scale_type.getRank() != 0 && scale_type.getRank() != 1) {
      return failure();
    }
    SmallVector<float> scale_factors;
    scale_factors.reserve(filter_scale_attr.size());
    for (auto value : filter_scale_attr.getValues<APFloat>()) {

// Converts the condition for an IfOp/WhileOp to a boolean value.
Value ConvertConditionToBoolean(Operation* op, Value cond) {
  if (auto ranked_type = mlir::dyn_cast<RankedTensorType>(cond.getType()))
    if (ranked_type.getRank() == 0 &&
        ranked_type.getElementType().isSignlessInteger(1))
      return cond;

  OpBuilder builder(op);
  Value to_bool = builder.create<TF::ToBoolOp>(op->getLoc(), cond);

  "GetHLOAxisFromTFAxis("
  "$0, $1.getType().cast<RankedTensorType>().getRank(), &$_builder)">;

// Same as the above but with $1 of type operand_range from variadic TensorFlow
// input.
def GetHLOAxisFromTFAxisVariadic : NativeCodeCall<
  "GetHLOAxisFromTFAxis("
  "$0, (*$1.begin()).getType().cast<RankedTensorType>().getRank(), "
  "&$_builder)">;

def CastElementsToI64Elements : NativeCodeCall<

                         (TFL_TopKV2Op $input, $k)>;

def ReductionDimensionIsLastDim : Constraint<CPred<"($0.cast<IntegerAttr>().getInt() == "
  "$1.getType().cast<ShapedType>().getRank() - 1 || $0.cast<IntegerAttr>().getInt() == -1)">>;

// Legalizes TF_ApproxTopKOp to TFL_TopKV2Op with the following constraints:
//    1. It computes max k
//    2. The reduction dimension is the last dim of the input.

// `ShapedType` or its rank is unknown.
inline bool HasRankOf(Value value, const int64_t rank) {
  auto shaped_type = mlir::dyn_cast_or_null<ShapedType>(value.getType());
  return shaped_type && shaped_type.hasRank() && shaped_type.getRank() == rank;
}

// Creates a new type that has the shape from the `old_type` and the element
// type from the `element_type`.
Type CloneTypeWithNewElementType(Type old_type, Type element_type);

        return (constant.getType().getRank() == 2);
      };

  auto op = cast<BatchMatMulOpType>(bmm_op);

  // Create a tfl.transpose op that performs ZX transpose on `input`.
  auto create_z_x_transpose_op = [&](Value input) -> Value {
    RankedTensorType input_type = mlir::cast<RankedTensorType>(input.getType());
    const int input_rank = input_type.getRank();

    return op.emitError("lhs must have static shape.");
  }
  if (!rhs_shape.hasStaticShape()) {
    return op.emitError("rhs must have static shape.");
  }

  const int64_t padding_nums_size = 2 * (rhs_shape.getRank() - 2);
  padding_nums.reserve(padding_nums_size);
  if (conv_padding.strref() == "EXPLICIT") {
    for (auto padding_elem :
         op.getExplicitPaddingAttr().template getAsRange<IntegerAttr>()) {

  }
  return success();
}

inline ShapedType MergeType(ShapedType a, ShapedType b) {
  if (!a.hasRank()) {
    return b;
  }
  if (!b.hasRank()) {
    return a;
  }
  int64_t rank = a.getRank();
  SmallVector<int64_t, 4> dims;
  dims.resize(rank);
  for (int i = 0, e = rank; i != e; i++) {
    int64_t dim0 = a.getDimSize(i);
    int64_t dim1 = b.getDimSize(i);