int num_send_from_host = 0, num_recv_at_host = 0;
  std::vector<Node *> send_recv_nodes;
  for (Node *n : host_graph->nodes()) {
    if (n->type_string() == "_XlaSendFromHost") {
      num_send_from_host++;
      send_recv_nodes.push_back(n);
    } else if (n->type_string() == "_XlaRecvAtHost") {
      num_recv_at_host++;
      send_recv_nodes.push_back(n);
    }
  }
  EXPECT_EQ(num_send_from_host, 1);

                        const std::vector<Output>& grad_inputs,
                        std::vector<Output>* grad_outputs) {
  ops::GradFunc grad_fn;
  TF_RETURN_IF_ERROR(ops::GradOpRegistry::Global()->Lookup(
      op.node()->type_string(), &grad_fn));
  TF_RETURN_IF_ERROR(grad_fn(scope, op, grad_inputs, grad_outputs));
  TF_RETURN_IF_ERROR(scope.status());
  return absl::OkStatus();
}

}  // end namespace test

//
// Unstage -> Node_Y
//
Status ClusterScopingPassImpl::ScopingForPipelineStages() {
  for (Node* n : graph_->nodes()) {
    DCHECK(n);
    if (n->type_string() == "Unstage") {
      AddScopeToAllTransitiveSuccessors(n);
    }
    if (n->type_string() == "Stage") {
      AddScopeToAllTransitivePredecessors(n);
    }
  }

  return absl::OkStatus();
}

Status ClusterScopingPassImpl::Run() {

    const OpRegistrationData* op_reg_data;
    TF_RETURN_IF_ERROR(library->LookUp(node->type_string(), &op_reg_data));
    if (op_reg_data->is_function_op) {
      function_nodes.push_back(node);
    }
  }
  for (auto node : function_nodes) {
    VLOG(2) << "Inlining function " << node->name();
    const FunctionDef* fdef = library->Find(node->type_string());
    if (fdef == nullptr) {

  if (pjrt_c_api_client == nullptr) {
    return absl::InternalError(absl::StrCat("PjRtClient for ",
                                            device_type.type_string(),
                                            " is not type PjRtCApiClient"));
  }
  return pjrt_c_api_client;
}

absl::Status ResetPjRtClient(const DeviceType& device_type) {

XlaDeviceDummyOp::XlaDeviceDummyOp(OpKernelConstruction* ctx) : OpKernel(ctx) {}

void XlaDeviceDummyOp::Compute(OpKernelContext* ctx) {
  LOG(FATAL) << "Attempted to execute Op " << name() << " type "
             << type_string() << " on an XLA device. This should never happen.";
}

XlaAssignVariableOp::XlaAssignVariableOp(OpKernelConstruction* c)
    : OpKernel(c) {
  OP_REQUIRES_OK(c, c->GetAttr("dtype", &dtype_));
}

          << op_kernel->type_string();
  ShowXlaDeviceDeprecationWarning(jit_device_name_.type_string());
  op_kernel->Compute(context);
}

void XlaDevice::ComputeAsync(AsyncOpKernel* op_kernel, OpKernelContext* context,
                             AsyncOpKernel::DoneCallback done) {
  ShowXlaDeviceDeprecationWarning(jit_device_name_.type_string());

  };

  if (node.type_string() == "IteratorGetNext" ||
      node.type_string() == "IteratorGetNextSync" ||
      node.type_string() == "MultiDeviceIteratorGetNextFromShard")
    return type_from_array_attr("output_shapes", "output_types");

  if (node.type_string() == "InfeedDequeueTuple")
    return type_from_array_attr("shapes", "dtypes");

  if (node.type_string() == "InfeedDequeue") {

    const XlaResourceOpInfo* op_info =
        GetResourceOpInfoForOp(node.type_string());
    return op_info && op_info->resource_kind() == XlaResourceKind::kStack;
  }

  bool IsTensorArrayOp(const Node& node) const {
    const XlaResourceOpInfo* op_info =
        GetResourceOpInfoForOp(node.type_string());
    return op_info && op_info->resource_kind() == XlaResourceKind::kTensorArray;
  }

    return n->type_string() == "NoOp" && n->out_edges().size() == 1 &&
           (*n->out_edges().begin())->dst()->IsSink();
  }

  return false;
}

bool MarkForCompilationPassImpl::IsScalarIntegerResourceOperation(
    const Cluster& cluster) {
  Node* n = GetOnlyNodeIn(cluster);
  if (!n) {
    return false;
  }

  if (n->type_string() != "AssignAddVariableOp" &&