def "does not reserve idle client when no clients"() {
        expect:
        manager.reserveIdleClient(options) == null
    }

    def "does not reserve idle client when no matching client found"() {
        def noMatch = Stub(WorkerDaemonClient) {
            isCompatibleWith(_) >> false
        }

  mlir::DictionaryAttr tf_attrs =
      entry_function->getAttrOfType<mlir::DictionaryAttr>("tf.entry_function");
  llvm::SmallVector<llvm::StringRef, 4> function_input_names;
  function_input_names.reserve(model_flags.input_arrays().size());
  auto input_attr = tf_attrs.get("inputs");
  if (!input_attr) {
    return errors::InvalidArgument("no inputs attribute found");
  }

                     llvm::ArrayRef<Value> outputs, OpBuilder* builder) {
  llvm::SmallVector<Type, 4> operand_types;
  operand_types.reserve(inputs.size());
  for (Value v : inputs) operand_types.emplace_back(v.getType());
  llvm::SmallVector<Type, 4> output_types;
  output_types.reserve(outputs.size());
  for (Value v : outputs) output_types.emplace_back(v.getType());

  mlir::OpBuilder builder(func.getBody());
  auto& bb = func.front();
  auto loc = func.getLoc();

  int num_args = bb.getNumArguments();
  llvm::SmallVector<Type, 4> input_types;
  input_types.reserve(num_args);
  // Edit the block arguments and create the new input ops in place to replace
  // the old input ops and quantize ops.
  for (int i = 0; i != num_args; ++i) {

  if (!attr) return success();

  // ID pairs are separated by `,`.
  llvm::SmallVector<std::string, 8> str_list =
      absl::StrSplit(attr.getValue().str(), ',', absl::SkipWhitespace());
  id_pairs.reserve(str_list.size());
  for (const std::string& str : str_list) {
    // IDs of one pair are separated by `:`.
    llvm::SmallVector<std::string, 8> id_pair = absl::StrSplit(str, ':');

    // Check for malformed IDs.

GetTPUDevices(ParsedDevices devices,
              llvm::ArrayRef<ParsedDevice> system_devices) {
  llvm::SmallVector<llvm::SmallVector<ParsedDevice, 8>, 8> tpu_devices;
  tpu_devices.reserve(system_devices.size());

  auto lookup = [&devices](ParsedDevice device_spec) {
    device_spec.has_type = true;
    device_spec.type = kDeviceTPU;
    // Enumerate all the available TPUs.

  std::vector<int64_t> targets;
  std::vector<Gradient*> used_in_grads;
  // We may end up with slightly fewer elements than we reserve, but grad.size()
  // should be a reasonably tight upper bound.
  targets.reserve(grad.size());
  used_in_grads.reserve(grad.size());
  std::unordered_map<int64_t, TapeTensor> sources_that_are_targets;
  for (int grad_index = 0, end = grad.size(); grad_index < end; ++grad_index) {

          IsWeightOnlyQuantizableOp(*candidate_op);

      // Collect all the quantized inputs and "clone" the matched op by these
      // inputs.
      SmallVector<Value, 4> inputs;
      inputs.reserve(candidate_op->getNumOperands());
      for (auto operand : candidate_op->getOperands()) {
        Type operand_type = operand.getType();
        if (mlir::isa<NoneType>(operand_type)) {
          inputs.push_back(operand);

    operand_new_shape[dimensions[i]] = operand.getType().getDimSize(i);
  }

  llvm::SmallVector<ElementValueT, 16> new_values;
  auto num_elements = result_type.getNumElements();
  new_values.reserve(num_elements);
  auto operand_values = operand.getValues<ElementValueT>();
  for (int64_t i = 0; i < num_elements; ++i) {
    const int64_t operand_index =
        GetElementIndex(operand_new_shape, current_index);

#include "xla/client/local_client.h"

namespace tensorflow {
namespace {
std::vector<const xla::Shape*> GetShapePointers(
    absl::Span<const xla::Shape> shapes) {
  std::vector<const xla::Shape*> shape_ptrs;
  shape_ptrs.reserve(shapes.size());
  for (const auto& shape : shapes) {
    shape_ptrs.push_back(&shape);
  }
  return shape_ptrs;
}
}  // namespace

absl::StatusOr<std::unique_ptr<xla::LocalExecutable>>