Clean up the build_xla_ops to use the generated C++ TF op wrappers.

This cleanup will make the future CL implementing lazy compilation simpler.

Includes some supporting changes:

 - Teach NewInternalScope to create a scope that doesn't do shape inference.  We
   need this because we don't have a ShapeRefiner that has been run over the
   entire graph available in the build_xla_ops pass.

 - Add a WithAssignedDevice modifier to tensorflow::Scope.

 - Make cc_op_gen write out an Operation field for nodes which may not
   necessarily have any outputs.  We already did this in most cases, but we
   weren't doing it for nodes that have possibly-empty list outputs.

 - Minor change renaming ops/xla_jit_op.cc to ops/xla_jit_ops.cc, now that we
   have more than one XLA JIT op.

PiperOrigin-RevId: 215293817

1 files changed, 5 insertions, 0 deletions

diff --git a/tensorflow/cc/framework/scope_internal.h b/tensorflow/cc/framework/scope_internal.h
index 58adaef2e9..514e02e841 100644
--- a/tensorflow/cc/framework/scope_internal.h
+++ b/tensorflow/cc/framework/scope_internal.h
@@ -26,6 +26,8 @@ class ShapeRefiner;
 // graph, status, name_map, and refiner.
 // This is intended to enable the C API (which are used by other language
 // bindings) to create a Scope and access C++ functionality (i.e. gradients).
+//
+// Shape inference is disabled if `refiner` is nullptr.
 Scope NewInternalScope(Graph* graph, Status* status, ShapeRefiner* refiner);
 
 class Scope::Impl {
@@ -58,6 +60,7 @@ class Scope::Impl {
     enum class ExitOnError;
     enum class KernelLabel;
     enum class Colocate;
+    enum class AssignedDevice;
   };
 
   Impl(Graph* graph, Status* status, NameMap* name_map, ShapeRefiner* refiner,
@@ -74,6 +77,7 @@ class Scope::Impl {
   Impl(const Scope& other, Tags::KernelLabel, const string& kernel_label);
   Impl(const Scope& other, Tags::Colocate, const Operation& colocate_with_op,
        bool clear_colocations);
+  Impl(const Scope& other, Tags::AssignedDevice, const string& assigned_device);
 
   std::unordered_set<string> GetColocationConstraints(
       const Operation& colocate_with_op) const;
@@ -107,6 +111,7 @@ class Scope::Impl {
   const bool exit_on_error_ = false;
   const string kernel_label_ = "";
   const string device_ = "";
+  const string assigned_device_ = "";
   const std::unordered_set<string> colocation_constraints_;
 
   // If true, Scope::DoShapeInference() always returns Status:OK().