// Here we assume that the element_shape won't be changed before calling
    // the first `TensorListSetItemOp`.
    if (auto shaped_type = element_shape.getType().dyn_cast<ShapedType>()) {
      if (shaped_type.hasRank() && shaped_type.getRank() == 0) {
        bool element_shape_acquired = false;
        auto uses = op.getResult().getUses();
        for (auto &use : llvm::make_early_inc_range(uses)) {

        TypedAttr typed_attr = it.dyn_cast<TypedAttr>();
        if (!typed_attr) return failure();
        all_types.insert(typed_attr.getType());
      }
      if (all_types.size() != 1) return failure();
      ShapedType new_out_type = RankedTensorType::get(
          {static_cast<int64_t>(array.size())}, *all_types.begin());
      DenseElementsAttr attr =
          DenseElementsAttr::get(new_out_type, array.getValue());

  int64_t prod_shape = 1;
  for (auto index_tuple : index_tuples) {
    const int64_t shape_i = shape[std::get<I>(index_tuple)];
    if (ShapedType::isDynamic(shape_i)) return ShapedType::kDynamic;
    prod_shape *= shape_i;
  }
  return prod_shape;
}

// Reshapes LHS and RHS to have B0,...,Bn,L,C and B0,...,Bn,C,R shape

    const auto [input_height, input_width] =
        GetDimSize(op->getOperand(0).getType().cast<ShapedType>().getShape(),
                   dimension_numbers.getInputSpatialDimensions());
    const auto [output_height, output_width] =
        GetDimSize(op->getResult(0).getType().cast<ShapedType>().getShape(),
                   dimension_numbers.getOutputSpatialDimensions());

  SmallVector<Attribute> shape_attrs;
  for (const Type result_type : output_types) {
    shape_attrs.push_back(
        tf_type::ShapeAttr::get(ctx, mlir::cast<ShapedType>(result_type)));
  }
  auto empty_array_attr = ArrayAttr::get(ctx, {});
  auto platforms = ArrayAttr::get(ctx, {StringAttr::get(ctx, kPlatformCpu)});

  auto call_op = builder.create<TF::XlaCallModuleOp>(

  // Create attributes used for creating an XlaCallModuleOp.
  SmallVector<Attribute> shape_attrs;
  for (const Type result_type : result_types) {
    shape_attrs.push_back(
        tf_type::ShapeAttr::get(ctx, mlir::cast<ShapedType>(result_type)));
  }
  const auto empty_array_attr = ArrayAttr::get(ctx, {});
  // TODO: b/310291615 - find a better way for platform support.
  const auto platforms = ArrayAttr::get(
      ctx,

                         (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;
  }

  if (!tensor.quantization || tfl::IsQuantized(tensor)) return nullptr;
  // If the result isn't float and unquantizable, the min/max is ignored.
  if (!res.getType()
           .cast<mlir::ShapedType>()
           .getElementType()
           .isa<mlir::FloatType>()) {
    return nullptr;
  }
  auto mins = tensor.quantization->min;
  auto maxs = tensor.quantization->max;

    // `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,