if (!mlir::isa<FloatType>(layerStatsType.getElementType())) {
      return emitOpError("layerStats must have a floating point element type");
    }
    if (layerStatsType.getRank() != 1 || layerStatsType.getDimSize(0) != 2) {
      return emitOpError("layerStats must have shape [2]");
    }
  }
  // Verify axisStats (optional) attribute.
  if (getAxisStats()) {

    if (!mlir::isa<FloatType>(layerStatsType.getElementType())) {
      return emitOpError("layerStats must have a floating point element type");
    }
    if (layerStatsType.getRank() != 1 || layerStatsType.getDimSize(0) != 2) {
      return emitOpError("layerStats must have shape [2]");
    }
  }
  // Verify axisStats (optional) attribute.
  if (getAxisStats()) {

  auto lhs_type = mlir::cast<ShapedType>(lhs.getType());
  auto rhs_type = mlir::cast<ShapedType>(rhs.getType());
  const int lhs_rank = lhs_type.getRank();
  const int rhs_rank = rhs_type.getRank();
  ImplicitLocOpBuilder builder(loc, rewriter);

  // Collects lhs and rhs dimensions information.
  DotDimensionsInfo lhs_dot_dimensions_info(

                   << " with broadcast_dims = " << broadcast_dims;
    return nullptr;
  }
  auto larger_broadcast_dims =
      GetI64ElementsAttrForSeq(0, result_type.getRank(), &builder);
  if (x_type.getRank() < y_type.getRank()) {
    if (x_type != result_type) {
      x = builder.create<BroadcastInDimOp>(loc, result_type, x, broadcast_dims);
    }
    if (y_type != result_type) {

    if (!type || !type.getElementType().isF32()) {
      return failure();
    }
    return success(
        op->hasOneUse() &&
        IsWeightQuantizableFunction(*op->getUses().begin(), type.getRank()));
  }

  // Checks if the operand is second operand of `tf.XlaCallModule` op for
  // `stablehlo.convolution` or `stablehlo.dot_general` with fully_quantizable
  // trait.

  // If the original type doesn't have a rank, then refine as the updated type
  // has a rank.
  if (!original.hasRank()) return true;

  // Both types must have the same rank.
  if (original.getRank() != updated.getRank()) return false;

  // Refine if the updated type is bounded.
  return IsBounded(updated);
}

// Propagates more refined type by cloning op using the new operands. This

  }

  // Currently we only support compressing weights of ops:
  //   Conv, DepthwiseConv, TransposeConv, whose filter has rank 4, and
  //   FullyConnected, whose filter has rank 2.
  if (type.getRank() != 2 && type.getRank() != 4) {
    result.can_compress = false;
    return result;
  }

  float random_sparsity = CalculateRandomSparsity(attr, type);
  if (random_sparsity < kMinSparsityLevel) {

  // Output type is unranked if input type is not ranked.
  auto ranked_ty = mlir::dyn_cast<RankedTensorType>(input_ty);
  if (!ranked_ty) return UnrankedTensorType::get(element_ty);
  int64_t rank = ranked_ty.getRank();

  DenseIntElementsAttr indices;
  if (!matchPattern(reduction_indices, m_Constant(&indices))) {
    // Output type is unranked if reduction indices are not constant and reduced
    // dimensions are not kept.

inline Type TransposeLastTwoDims(Type type) {
  auto shaped_type = type.dyn_cast<ShapedType>();
  if (!shaped_type.hasStaticShape() || shaped_type.getRank() < 2) {
    return nullptr;
  }
  int rank = shaped_type.getRank();
  if (rank < 2) {
    return nullptr;
  }
  SmallVector<int64_t> new_shape(shaped_type.getShape().begin(),
                                 shaped_type.getShape().end());

  auto input_type = mlir::cast<RankedTensorType>(input.getType());
  SmallVector<int64_t, 4> output_shape;
  int size_of_splits;
  if (input_type.getRank() < axis || axis < 0) return failure();
  for (int i = 0; i < input_type.getRank(); ++i) {
    int64_t dim = input_type.getDimSize(i);
    if (i == axis) {
      if (dim % splits != 0) {
        return failure();
      }