Searched refs:on_device (Results 1 – 8 of 8) sorted by relevance
| /external/tensorflow/tensorflow/core/profiler/convert/ |
| D | op_stats_to_tf_stats_test.cc | 37 bool on_device, in AddTensorFlowOpEvent() argument
40 absl::string_view name = on_device ? kernel_name : tf_op_fullname; in AddTensorFlowOpEvent()
44 if (!on_device) return event; in AddTensorFlowOpEvent()
53 int64 duration_ns, bool on_device, in AddTensorFlowOpEventWithKernelDetails() argument
60 duration_ns, on_device, kernel_name, plane, line); in AddTensorFlowOpEventWithKernelDetails()
61 if (!on_device) return; in AddTensorFlowOpEventWithKernelDetails()
|
| D | xplane_to_op_metrics_db_test.cc | 36 bool on_device, absl::string_view kernel_name, in AddTensorFlowOpEvent() argument
38 absl::string_view name = on_device ? kernel_name : tf_op_fullname; in AddTensorFlowOpEvent()
42 if (!on_device) return; in AddTensorFlowOpEvent()
|
| D | op_stats_to_tf_stats.cc | 37 bool on_device, const OpMetrics& metrics, in ConvertOpMetricsToTfStatsRecord() argument
40 record.set_host_or_device(on_device ? "Device" : "Host"); in ConvertOpMetricsToTfStatsRecord()
|
| /external/tensorflow/tensorflow/compiler/xla/python/ |
| D | jax_jit.cc | 1009 TF_ASSIGN_OR_RETURN(DevicePutResult on_device, in ConvertArgsToBuffers()
1012 xla::PjRtBuffer* buffer = on_device.buffer; in ConvertArgsToBuffers()
1014 if (on_device.owned_buffer) { in ConvertArgsToBuffers()
1015 keep_alive.emplace_back(std::move(on_device.owned_buffer)); in ConvertArgsToBuffers()
1020 on_device.weak_type); in ConvertArgsToBuffers()
|
| /external/tensorflow/tensorflow/python/distribute/ |
| D | packed_distributed_variable.py | 87 def on_device(self, device): member in PackedDistributedVariable
|
| D | tpu_strategy.py | 1178 value._packed_variable.on_device(d)
1241 value._packed_variable.on_device(d)
|
| D | values.py | 682 return self._packed_var.on_device(self._devices[replica_id])
|
| /external/tensorflow/tensorflow/compiler/xla/service/spmd/ |
| D | spmd_partitioner.cc | 2192 auto on_device = b_.AddInstruction( in HandleSingleDevice() local
2195 ShapeUtil::MakeShape(PRED, {}), partition_id_, on_device, in HandleSingleDevice()
|