(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

    // unified_api_test.cc.
    return Unimplemented("MlirTensor::Shape is not implemented yet.");
  }

  Value getValue() { return value_; }
  Type getElementType() {
    return mlir::cast<ShapedType>(value_.getType()).getElementType();
  }

  // For LLVM style RTTI.
  static bool classof(const AbstractTensorHandle* ptr) {
    return ptr->getKind() == kMlir;
  }

    // `stablehlo.dot_general` legalizable to `tfl.fully_connected` has a
    // filter rank of 2 with the last dimension as the channel dimension.
    const int64_t quantization_dimension =
        mlir::cast<ShapedType>(filter_type).getShape().size() - 1;
    accumulation_quantized_element_type =
        CreateI32F32UniformQuantizedPerAxisType(
            gemm_style_op->getLoc(), *rewriter.getContext(), result_scales,

      ranked_ty.getEncoding());
  if (encoding && !encoding.getBounds().empty()) {
    for (int64_t dim = 0; dim < rank; ++dim) {
      if (dims[dim] == mlir::ShapedType::kDynamic) {
        dims[dim] = encoding.getBounds()[dim];
      }
    }
  }
  return GetTypeFromTFTensorShape(dims, ranked_ty.getElementType());
}