return std::nullopt;  // `call_op` is not a quantized function call.
  }

  absl::StatusOr<Method> method = GetQuantizationMethod(call_op);
  if (!method.ok()) {
    call_op->emitError() << "Failed to get quantization method: "
                         << method.status().ToString();
    return std::nullopt;
  }

  QuantizationResult result{};
  result.mutable_quantizable_unit()->set_name(

  // module level.
  for (auto func : module_op.getOps<func::FuncOp>()) {
    if (failed(applyPatternsAndFoldGreedily(func, frozen_patterns))) {
      func.emitError() << "quant-propagate-quantize-type failed.";
      signalPassFailure();
    }
  }
}

}  // namespace

// Creates an instance of the TensorFlow dialect PropagateQuantizeType pass.

      fetched_device_coordinates.ok()) {
    device_coordinates = *fetched_device_coordinates;
    return mlir::success();
  } else {
    return cluster.emitError() << "error in fetching tpu device coordinates: "
                               << fetched_device_coordinates.status().message();
  }
}

int GetNumCoresPerReplica(mlir::tf_device::ClusterOp cluster) {

    }

    Value input = tf_op.getInputs();
    int quant_dim = -1;
    auto input_type = mlir::cast<ShapedType>(input.getType());
    if (PerAxis) {
      if (!input_type.hasRank()) {
        tf_op.emitError("The input should have known rank for per-channel op.");
        return failure();
      }
      // This is a special case that the quant_dim is the last dimensions.
      quant_dim = input_type.getRank() - 1;
    }

      << ": " << num_part_ops << " ops delegated out of " << num_all_ops
      << " ops with " << num_partitions_after
      << " partitions (originally: " << num_partitions_before << ")";
  if (!scheduled_everything) {
    emitError(func.getLoc(), func.getName())
        << ": " << num_unscheduled_ops << " operations couldn't be scheduled";
  }
  return scheduled_everything;
}

class PartitionedTopologicalSortPass

                                 &send_gradient_op_map)))
    return signalPassFailure();

  if (enqueue_op_map.size() != recv_activation_op_map.size()) {
    func_op.emitError() << "expects the number of embedding enqueue ops to "
                           "match the number of '"
                        << TF::RecvTPUEmbeddingActivationsOp::getOperationName()
                        << "' ops";

  MLIRContext& ctx = getContext();

  RewritePatternSet patterns(&ctx);
  patterns.add<RewriteNchwConvolutionToNhwc>(&ctx);

  if (failed(applyPatternsAndFoldGreedily(func_op, std::move(patterns)))) {
    func_op.emitError() << "Failed to run NchwConvolutionToNhwcPass.";
    signalPassFailure();
  }
}

  RewritePatternSet patterns(ctx);
  func::FuncOp func = getOperation();

  patterns.add<AddQuantizationUnitLoc>(ctx);
  if (failed(applyPatternsAndFoldGreedily(func, std::move(patterns)))) {
    func.emitError() << "quant-add-quantization-unit-loc pattern "
                        "conversion did not converge.";
    signalPassFailure();
  }
}

}  // namespace

  // Handle only static shape cases.
  // TODO(b/260284866): Handle dynamic shape cases.
  if (!lhs_shape.hasStaticShape()) {
    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);

  tf_executor::GraphOp graph_op;

  if (llvm::hasSingleElement(graph_op_range))
    graph_op = dyn_cast<tf_executor::GraphOp>(func.front().front());

  if (!graph_op) {
    func.emitError("expected function to contain only a graph_op");
    signalPassFailure();
    return;
  }

  // New control inputs to be added. For an operation x, new_control_inputs[x]