namespace tensorflow {

// Checks if a graph or reachable functions in the library have any
// StatefulPartitionedOps with _XlaMustCompile=true. The function library will
// be skipped if nullptr is provided.
bool IsSupportedByNonReplicatedBridge(
    const Graph& graph, const FunctionLibraryDefinition* function_library);

// Checks if a graph or reachable functions in the library have any ops with

      /*arg_def=*/{"in: resource"},
      /*ret_def=*/{"out: resource"},
      /*attr_def=*/{},
      /*node_def=*/
      {{{"out"}, "Identity", {"in"}, {{"T", DataType::DT_RESOURCE}}}});
}

TEST(IsSupportedByNonReplicatedBridge, NonReplicatedGraph) {
  const FunctionDef& fd = PassThroughResource();
  FunctionDefLibrary flib;
  *flib.add_function() = fd;
  FunctionLibraryDefinition flib_def(OpRegistry::Global(), flib);

    for (const NodeDef& node_def : func_def->node_def()) {
      if (node_def.op() == kPartitionedCall) return true;
    }
  }
  return false;
}

}  // namespace

bool IsSupportedByNonReplicatedBridge(
    const Graph& graph, const FunctionLibraryDefinition* function_library) {
  return IsNonReplicatedGraph(graph, function_library) &&
         HasPsWithResourceVariable(graph);
}