A closure describing how to run a compiled version of a TensorFlow function. It may seem unusual to stick the resource variable snapshots in this class. This is necessary: we need to use the snapshots observed by the compiler as the initial values for the resource variables (and cannot snapshot them again during execution) because otherwise we risk observing a different snapshot with shapes diffe
| 101 | // during execution) because otherwise we risk observing a different snapshot |
| 102 | // with shapes different from what we compiled for. |
| 103 | class XlaExecutableClosure { |
| 104 | public: |
| 105 | explicit XlaExecutableClosure( |
| 106 | xla::LocalClient* client, xla::LocalExecutable* executable, |
| 107 | const XlaCompiler::CompilationResult* compilation_result, |
| 108 | std::map<int, OptionalTensor> resource_var_snapshots, |
| 109 | int num_constant_args) |
| 110 | : client_(client), |
| 111 | executable_(executable), |
| 112 | compilation_result_(compilation_result), |
| 113 | resource_var_snapshots_(std::move(resource_var_snapshots)), |
| 114 | num_constant_args_(num_constant_args) {} |
| 115 | |
| 116 | XlaExecutableClosure(XlaExecutableClosure&&) = default; |
| 117 | XlaExecutableClosure& operator=(XlaExecutableClosure&&) = default; |
| 118 | |
| 119 | xla::LocalClient* client() const { return client_; } |
| 120 | xla::LocalExecutable* executable() const { return executable_; } |
| 121 | const XlaCompiler::CompilationResult* compilation_result() const { |
| 122 | return compilation_result_; |
| 123 | } |
| 124 | const std::map<int, OptionalTensor>& resource_var_snapshots() const { |
| 125 | return resource_var_snapshots_; |
| 126 | } |
| 127 | int num_constant_args() const { return num_constant_args_; } |
| 128 | |
| 129 | private: |
| 130 | xla::LocalClient* client_; |
| 131 | xla::LocalExecutable* executable_; |
| 132 | const XlaCompiler::CompilationResult* compilation_result_; |
| 133 | std::map<int, OptionalTensor> resource_var_snapshots_; |
| 134 | int num_constant_args_; |
| 135 | |
| 136 | TF_DISALLOW_COPY_AND_ASSIGN(XlaExecutableClosure); |
| 137 | }; |
| 138 | |
| 139 | // This maintains a mapping from a globally unique ID to XlaExecutableClosure |
| 140 | // instances. |