ConvertAssetArgsTest() {
    ctx_.loadDialect<func::FuncDialect, TF::TensorFlowDialect,
                     tf_saved_model::TensorFlowSavedModelDialect>();
  }

  // Parses `module_op_str` to create a `ModuleOp`. Checks whether the created
  // module op is valid.
  OwningOpRef<ModuleOp> ParseModuleOpString(
      const absl::string_view module_op_str) {
    auto module_op_ref = parseSourceString<ModuleOp>(module_op_str, &ctx_);

  FunctionType function_type =
      builder.getFunctionType(input_types, return_types);

  std::string function_name = absl::StrCat("func_", subgraph.subgraph_id_);

  func::FuncOp new_func = func::FuncOp::create(builder.getUnknownLoc(),
                                               function_name, function_type);
  new_func.setVisibility(func::FuncOp::Visibility::Private);
  new_func.addEntryBlock();

  auto var_handle_op_in_init = var_handle_op->clone();
  builder.insert(var_handle_op_in_init);
  auto const_op = builder.create<mlir::arith::ConstantOp>(
      session_init_func.getLoc(), tensor_attr.getType(), tensor_attr);

  builder.create<TF::AssignVariableOp>(
      session_init_func.getLoc(), llvm::ArrayRef<mlir::Type>{},
      llvm::ArrayRef<mlir::Value>{var_handle_op_in_init->getResult(0),

  }

  if constexpr (is_input) {
    auto pi = builder.create<TF::TPUPartitionedInputV2Op>(
        op.getLoc(), op.getType(), op.getOperands(),
        builder.getI64ArrayAttr(partition_dims), builder.getBoolAttr(false),
        op.get_XlaShardingAttr());
    op->replaceAllUsesWith(pi);
  } else {
    auto po = builder.create<TF::TPUPartitionedOutputV2Op>(
        op.getLoc(), op.getResultTypes(), op.getOperand(),

  } else {
    return false;
  }
  return true;
}

// Create a dummy zero to be fed locally from the host to the TPUExecute.
Value CreateZeroInput(Location loc, OpBuilder& builder, Attribute zero_attr,
                      DenseIntElementsAttr shape_attr) {
  ConstOp zero = builder.create<ConstOp>(loc, zero_attr);
  zero->setAttr(kICIWeightDistributionMlirBridgeMarker,

    config.mutable_calibration_options()
        ->mutable_representative_datasets()
        ->Add(preset_datasets.begin(), preset_datasets.end());
  }

  // Create a new `QuantizationSpecs` to replace the existing one. The
  // expansion from `StaticRangePtqPreset` gets populated first and then
  // user-provided explicit `QuantizationSpec`s will be appended.
  QuantizationSpecs new_specs{};

                                    /*allow_partial_conversion=*/false);

  auto pass_instrumentors = mlir::GetPassInstrumentors();
  for (const auto& creator : pass_instrumentors) {
    tf2xla.addInstrumentation(creator());
  }
  if (DEBUG_DATA_DUMPER()->ShouldDump(module_name.str(), kDebugGroupMain) ||
      VLOG_IS_ON(1)) {
    tensorflow::DumpMlirOpToFile(

    return convert(attr);
  }
  return nullptr;
}

DenseElementsAttr UniformQuantizedPerAxisValueConverter::convert(
    DenseFPElementsAttr attr) {
  // Creates the converter for each chunk. Normally the size of the
  // quantization dim is 3, so we can cache all the converters.
  ShapedType type = attr.getType();
  std::size_t dim_size = type.getDimSize(quantization_dim_);

  for (int i = 0; i < y_num; i++) {
    dy_datas[i] = Tensor(ys[i].type(), y_shapes[i]);
    auto dy_data_flat = dy_datas[i].flat<Y_T>();
    dy_data_flat.setZero();
  }

  // Create the feed list.
  ClientSession::FeedType feed_list;
  for (int i = 0; i < x_num; i++) {
    feed_list.insert({xs[i], x_datas[i]});
  }
  for (int i = 0; i < y_num; i++) {
    feed_list.insert({dys[i], dy_datas[i]});

    compilation_successful.scalar<bool>()() = false;
    ctx->set_output(0, compilation_key);
    ctx->set_output(1, compilation_successful);
    return;
  }

  // Each execution of an XlaCompile op creates a new ExecutableClosure, even
  // if it didn't have to compile the cluster because of a compilation-cache
  // hit.  This is because we at least need new snapshots of the resource
  // variables.