<< " 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) {

    // tensor, for setting depth_multiplier attribute, etc.).
    auto filter = tf_op.getFilter();
    auto filter_type = mlir::dyn_cast<RankedTensorType>(filter.getType());
    if (!filter_type || filter_type.getRank() != 4 ||
        !filter_type.hasStaticShape())
      return failure();

    Value input = tf_op.getInput();
    RankedTensorType input_type =
        mlir::dyn_cast<RankedTensorType>(input.getType());

    // before broadcasting.
    if (operand_type.getRank() < output_type.getRank()) {
      input = InsertExpandDimsOp(op, rewriter, input, output_type.getRank());
    }

    SmallVector<int32_t> broadcast_shape =
        CastI64ArrayToI32(output_type.getShape()).value();
    TensorType broadcast_shape_type =
        output_type.cloneWith({output_type.getRank()}, rewriter.getI32Type());
    auto broadcast_shape_attr =

  if (!image_type || !image_type.getElementType().isF32() ||
      image_type.getRank() != 4) {
    return func_.emitWarning() << "Image should be a 4D float tensor";
  }

  auto flow_type = mlir::dyn_cast_or_null<RankedTensorType>(
      func_.getFunctionType().getInput(1));
  if (!flow_type || !flow_type.getElementType().isF32() ||
      flow_type.getRank() != 4) {
    return func_.emitWarning() << "Flow should be a 4D float tensor";

        llvm::dyn_cast<RankedTensorType>(tpose_arg1->getResultTypes()[0]);
    auto tpose_arg2_type =
        llvm::dyn_cast<RankedTensorType>(tpose_arg2->getResultTypes()[0]);
    if (tpose_arg1_type.getRank() != tpose_arg2_type.getRank()) {
      return failure();
    }

    if (llvm::isa<BlockArgument>(tpose_arg1.getPerm()) ||
        llvm::isa<BlockArgument>(tpose_arg2.getPerm())) {
      return failure();
    }

  Block &cond = while_op.getCond().front();
  Block &body = while_op.getBody().front();

  auto n = while_op.getNumOperands();
  if (!AllOperationSafe(cond) || !AllOperationSafe(body)) return false;

  // Find all Consumed indices.
  // i is consumed element if result(i) is used outside whileOp or
  // argument(i) is used in whileOp.getCond().
  for (auto i = 0; i < n; ++i) {

      while_op->erase();
    } else if (auto while_op = dyn_cast<TF::WhileRegionOp>(user_op)) {
      auto new_while_op = GetUpdatedWhileOp(
          while_op, while_op.getCond().getArgumentTypes(), args_to_erase);
      new_while_op.getCond().takeBody(while_op.getCond());
      new_while_op.getBody().takeBody(while_op.getBody());
      llvm::BitVector erase_indices;
      UpdateTerminatorArguments(new_while_op.getBody(), args_to_erase,

    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.

  auto iota_shape = iota_type.getShape();
  auto reduce_dim = (*dimensions.value_begin<APInt>()).getSExtValue();
  if (reduce_dim < 0) reduce_dim += iota_type.getRank();

  auto index =
      std::optional<SmallVector<int64_t>>(std::in_place, iota_type.getRank());
  while (index.has_value()) {
    StridedArrayView<DenseIntElementsAttr> array_view(
        iota_const_attr, iota_shape, *index, reduce_dim);

    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()) {