Commong tracing function for both CPU and TPUs. The caller function should set device_type, num_replicas, num_replicas_per_host, num_hosts and replica_id before calling _trace_execution. Args: graph: the graph of Ops executed on the TPU. tensor_fetches: a (list,tuple,o
(self, graph,
tensor_fetches,
op_fetches=None,
on_tpu=True)
| 1264 | return self._host_call_fn |
| 1265 | |
| 1266 | def _trace_execution(self, graph, |
| 1267 | tensor_fetches, |
| 1268 | op_fetches=None, |
| 1269 | on_tpu=True): |
| 1270 | """Commong tracing function for both CPU and TPUs. |
| 1271 | |
| 1272 | The caller function should set device_type, num_replicas, |
| 1273 | num_replicas_per_host, num_hosts and replica_id before calling |
| 1274 | _trace_execution. |
| 1275 | |
| 1276 | |
| 1277 | Args: |
| 1278 | graph: the graph of Ops executed on the TPU. |
| 1279 | tensor_fetches: a (list,tuple,or a single object) of tensor fetches |
| 1280 | returned by model_fn given to session.run. Function must be provided |
| 1281 | with as least one tensor to fetch. |
| 1282 | op_fetches: A list of op fetches returned by model_fn given to |
| 1283 | session.run. op_fetches and tensor_fetches are used to determine the |
| 1284 | nodes that will be executed. Can be None. |
| 1285 | on_tpu: True if executing on TPU. |
| 1286 | |
| 1287 | Returns: |
| 1288 | tensor_fetches: an exact copy of tensor_fetches that has additional |
| 1289 | dependencies. |
| 1290 | Raises: |
| 1291 | RuntimeError: If tensor_fetches is None or empty. |
| 1292 | """ |
| 1293 | def _cast_unsupported_dtypes(tensor): |
| 1294 | """Casts tensor to a supported type.""" |
| 1295 | |
| 1296 | if tensor.dtype.__eq__(dtypes.int64): |
| 1297 | # outside-compilation doesn't support int64 input yet. |
| 1298 | return math_ops.cast(tensor, dtypes.int32) |
| 1299 | if tensor.dtype.__eq__(dtypes.bfloat16) or tensor.dtype.__eq__( |
| 1300 | dtypes.float16): |
| 1301 | # Since host can't handle bf16, convert tensor to f32. |
| 1302 | return math_ops.cast(tensor, dtypes.float32) |
| 1303 | return tensor |
| 1304 | |
| 1305 | TensorTracer.check_device_type(self._tt_config.device_type) |
| 1306 | TensorTracer.check_trace_mode(self._tt_config.device_type, |
| 1307 | self._parameters.trace_mode) |
| 1308 | # Check in_tensor_fetches, and op_fetches and convert them to lists. |
| 1309 | processed_t_fetches = self._process_tensor_fetches(tensor_fetches) |
| 1310 | op_fetches = self._process_op_fetches(op_fetches) |
| 1311 | all_fetches = op_fetches + [tensor.op for tensor in processed_t_fetches] |
| 1312 | |
| 1313 | # Filter out the operations that won't be executed. |
| 1314 | # if fetches=None, then ops_in_exec_path = set(operations) |
| 1315 | exec_op_set = self._filter_execution_path_operations(graph.get_operations(), |
| 1316 | all_fetches) |
| 1317 | # Write report file, and determine the traced tensors. |
| 1318 | tensor_trace_order = self._determine_trace_and_create_report( |
| 1319 | graph, exec_op_set) |
| 1320 | |
| 1321 | tensor_fetch_set = set(processed_t_fetches) |
| 1322 | tracing_ops = [] |
| 1323 |
no test coverage detected