limitations under the License.
==============================================================================*/

#include "tensorflow/compiler/mlir/tf2xla/internal/utils/test_metadata_config.h"

#include <vector>

#include "absl/status/status.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project

      source_mgr->getMemoryBuffer(source_mgr->getMainFileID());
  std::string error;
  auto loc =
      mlir::FileLineColLoc::get(context, input->getBufferIdentifier(), 0, 0);
  std::vector<std::string> inputs;
  std::vector<std::string> outputs;
  return tflite::FlatBufferToMlir(
      absl::string_view(input->getBufferStart(), input->getBufferSize()),
      context, loc, false, inputs, outputs, false);
}

}  // namespace

namespace tensorflow {
namespace generator {

PathConfig::PathConfig(const string& output_dir, const string& source_dir,
                       const string& api_dir_list,
                       const std::vector<string> op_names)
    : output_path(output_dir), op_names(op_names) {
  api_dirs = str_util::Split(api_dir_list, ",", str_util::SkipEmpty());

  // Decompose the directory components given the output/source directories.
  //

//   success, or an absl::Status with an error message on failure.
absl::StatusOr<OwningOpRef<ModuleOp>> TfToStablehlo(
    absl::string_view input_path, MLIRContext* context,
    const std::vector<std::string>& exported_model_signatures,
    const std::vector<std::string>& tag_names,
    absl::string_view input_arg_shapes_str, bool is_input_mlir_module);

}  // namespace mlir

#include "tensorflow/compiler/mlir/quantization/stablehlo/cc/calibration/component.h"

#include <optional>
#include <string>
#include <unordered_set>
#include <utility>
#include <vector>

#include "absl/algorithm/container.h"
#include "absl/base/nullability.h"
#include "absl/container/flat_hash_map.h"
#include "absl/log/die_if_null.h"
#include "absl/log/log.h"
#include "absl/status/status.h"

namespace tensorflow {
namespace generator {

string Call(const string& object, const string& method,
            std::vector<string> arguments, const char* oper) {
  return absl::Substitute("$0$1$2($3)", object, oper, method,
                          absl::StrJoin(arguments, ", "));
}

string Call(const string& function, std::vector<string> arguments) {
  return absl::Substitute("$0($1)", function, absl::StrJoin(arguments, ", "));
}

  }

  mlir::Location loc =
      mlir::FileLineColLoc::get(context, input_filename, /*line=*/0,
                                /*column=*/0);
  std::vector<std::string> inputs;
  std::vector<std::string> outputs;
  return tflite::FlatBufferToMlir(
      absl::string_view(buffer->getBufferStart(), buffer->getBufferSize()),
      context, loc, /*use_external_constant=*/false, inputs, outputs,

  mlir::MLIRContext context(registry);
  context.loadAllAvailableDialects();

  bool is_input_mlir_module = absl::EndsWith(input_path, ".mlir");
  std::vector<std::string> exported_model_signatures_in_vector =
      absl::StrSplit(exported_model_signatures, ',');
  std::vector<std::string> tag_names_in_vector = absl::StrSplit(tag_names, ',');
  auto module = TfToStablehlo(
      input_path, &context, exported_model_signatures_in_vector,

      tensorflow::io::JoinPath(*g_test_model_dir, internal::kModelWithCustomOp);
  return FlatBufferModel::BuildFromFile(model_path.c_str());
}

template <typename T>
std::vector<T> GetAsVector(const flatbuffers::Vector<T>* vec) {
  return std::vector<T>(vec->begin(), vec->end());
}

class QuantizeWeightsTest : public testing::Test {
 protected:
  QuantizeWeightsTest() {}

  void LoadBasicModel() {

class TpuNodeDeviceFactory : public DeviceFactory {
 public:
  Status ListPhysicalDevices(std::vector<string>* devices) override;
  Status CreateDevices(const SessionOptions& options, const string& name_prefix,
                       std::vector<std::unique_ptr<Device>>* devices) override;
};

Status TpuNodeDeviceFactory::ListPhysicalDevices(std::vector<string>* devices) {
  tpu::TpuPlatformInterface* platform =