}
  }
}

// Finds the QuantizationUnit from location.
std::optional<QuantizationUnit> FindQuantizationUnit(Operation* op) {
  SmallVector<QuantizationUnit> quant_units;
  FindQuantizationUnitsRecursively(op->getLoc(), quant_units);

  if (quant_units.size() == 1) {
    return *quant_units.begin();
  }
  // Among units, return the one with the same type as given op.
  StringRef given_op_type = op->getName().getStringRef();

      ++num_unscheduled_ops;
    }
  }
  emitRemark(func.getLoc(), func.getName())
      << ": " << 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())

    if (variable_name.empty()) {
      return signalPassFailure();
    }
    VarHandleOp var_handle_op = builder.create<VarHandleOp>(
        variable_v2_op.getLoc(),
        ArrayRef<Type>{RankedTensorType::get(
            {}, TF::ResourceType::get(ArrayRef<TensorType>{tensor_type},
                                      builder.getContext()))},
        ArrayRef<Value>{},

    // and its users.
    auto quantize = rewriter.create<quantfork::QuantizeCastOp>(
        tf_op.getLoc(), qtype.getValue(), input);
    auto dequantize = rewriter.create<quantfork::DequantizeCastOp>(
        tf_op.getLoc(), res_type, quantize.getResult());
    tf_op.getOutputs().replaceAllUsesWith(dequantize);

    return success();
  }

  bool use_fake_quant_num_bits_;

  if (IsKVCacheCompositeOp(composite)) {
    return builder.create<TFL::CustomOp>(
        composite->getLoc(), composite->getResultTypes(),
        composite->getOperands().slice(2, 3), composite.getName(),
        CustomOption(&builder, custom_option_buffer));
  }
  return builder.create<TFL::CustomOp>(
      composite->getLoc(), composite->getResultTypes(),
      composite->getOperands(), composite.getName(),

    // and its users.
    Value value = tf_op.getOutputs();
    auto quantize = rewriter.create<TFL::QuantizeOp>(
        tf_op.getLoc(), qtype.getValue(), value, qtype);
    auto dequantize = rewriter.create<TFL::DequantizeOp>(
        tf_op.getLoc(), res_type, quantize.getOutput());
    value.replaceAllUsesWith(dequantize);
    quantize.getOperation()->replaceUsesOfWith(dequantize, value);

      if (!dict_attr) continue;

      OpBuilder builder(func.getContext());
      builder.setInsertionPointToStart(&func.getBlocks().front());
      builder.create<TFL::CallOnceOp>(func.getLoc(), init_func_op.getName());
    }
  }
}

}  // namespace

// Inserts a TFL::CallOnce op when tf_saved_model's session initializer is
// given.
std::unique_ptr<OperationPass<ModuleOp>>

  auto restore_op = builder.create<TF::RestoreV2Op>(
      session_init_func.getLoc(),
      /*tensors=*/tensor_types,
      /*prefix=*/filename_arg, tensor_names_const, shape_and_slices_const);

  for (auto [idx, restore_result] : llvm::enumerate(restore_op.getResults())) {
    builder.create<TF::AssignVariableOp>(
        restore_op.getLoc(), target_var_handle_ops[idx], restore_result);
  }
}

  Block* block = builder->createBlock(body);

  // Block arguments are scalars of the given element type.
  Type type = RankedTensorType::get(/*shape=*/{}, element_type);
  Location loc = body->getLoc();
  block->addArguments({type, type}, SmallVector<Location, 2>(2, loc));

  auto reducer =
      builder->create<Op>(loc, block->getArgument(0), block->getArgument(1));
  builder->create<ReturnOp>(loc, reducer.getResult());

        transposer.TransposeValues(original_values);

    // Create a new constant op with the transposed values.
    const Location combined_loc =
        rewriter.getFusedLoc({const_op.getLoc(), op.getLoc()});
    auto new_value_type =
        RankedTensorType::getChecked(combined_loc, transposer.GetTargetShape(),
                                     /*elementType=*/rewriter.getF32Type());
    auto new_value_attr =