Type resultType = RankedTensorType::get(shape, element_type);
  auto constant_attr = DenseElementsAttr::get(shape_spec_type, shape);
  auto shape_tensor =
      rewriter.create<TF::ConstOp>(loc, shape_spec_type, constant_attr);
  return rewriter.create<TF::ReshapeOp>(loc, resultType, /*tensor=*/value,
                                        /*shape=*/shape_tensor);
}

template <typename BatchMatMulOpType>

  func->setAttr(kTFImplements, attr);
  OpBuilder builder(func.getBody());
  std::string empty_option_buffer;
  auto op = builder.create<CustomOp>(
      func.getLoc(), func.getFunctionType().getResults(), func.getArguments(),
      api, CustomOption(&builder, empty_option_buffer));
  builder.create<func::ReturnOp>(func.getLoc(), op.getResults());
  return success();
}

LogicalResult VerifyNgrams(func::FuncOp func) {

                                    /*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(

  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.

                          std::vector<Output>* grad_outputs);

  // Returns a list mapping whether each node in the graph is reachable
  // from outputs_. Keyed by node id.
  std::vector<bool> GetReachableNodes();

  // Creates the gradient subgraph for a while loop (or just stores
  // `summed_grads` if not all incoming gradients are available yet). All exit
  // nodes (which are the first nodes of a loop encountered in the backwards

  const std::string& profiler_name =
      GetPjRtDeviceCompilationProfilerResourceName(device_type);
  bool deleted_old_device_compiler = false;

  // Lookup the DeviceCompiler, create one if not found.
  Status s = rm->Lookup<PjRtDeviceCompiler>(
      rm->default_container(), compiler_name, pjrt_device_compiler);
  if (s.ok() && device_type == DEVICE_TPU) {

  for (const auto& device : devices)
    if (DeviceNameUtils::IsCompleteSpecification(spec, device))
      matching_devices.push_back(device);
  return matching_devices;
}

// Create error message for a conflicting attribute of a device.
template <typename T>
absl::Status MismatchedTPUSystemAttributeErr(absl::string_view attribute, T a,
                                             T b) {

    return a->isBeforeInBlock(b);
  });

  // Create tf_device::ClusterOp before the last operation in the block that
  // is a part of a match set.
  auto back = cluster.operations.back();
  auto loc = back->getLoc();
  OpBuilder builder(back);

  auto cluster_op =
      builder.create<tf_device::ClusterOp>(loc, return_types, policy);

  // Create block in cluster_op's region and move 'cluster.operations' into

    "bad '{0}' attribute at index {1} with value '{2}': failed to parse to {3}";
constexpr char kBadArrayAttrLengthMsg[] =
    "bad '{0}' attribute, expected array attribute of size {1}, got size {2}";

// Creates a missing attribute error message.
std::string CreateMissingAttributeMsg(llvm::StringRef attribute) {
  return llvm::formatv("requires attribute '{0}'", attribute).str();
}