| /external/tensorflow/tensorflow/python/tpu/ |
| D | tpu_context.py | 27 from tensorflow.python.tpu import device_assignment as tpu_device_assignment
135 def device_assignment(self): member in TPUContext
140 return self._internal_ctx.device_assignment
282 device_assignment = self._lazy_device_assignment_dict.get(master)
283 if device_assignment is not None:
284 return device_assignment
288 device_assignment = tpu_device_assignment.device_assignment(
298 str(device_assignment.topology.device_coordinates))
300 str(device_assignment.core_assignment))
302 self._lazy_device_assignment_dict[master] = device_assignment
[all …]
|
| D | tpu.py | 517 device_assignment=None, argument
573 device_assignment,
664 device_assignment=None, argument
722 if device_assignment is not None:
727 device_assignment.topology.serialized(),
729 device_assignment.core_assignment.flatten().tolist()
735 device_assignment.num_cores_per_replica)
738 device_assignment.num_cores_per_replica
1092 device_assignment=None, argument
1185 device_assignment=device_assignment,
[all …]
|
| D | tpu_feed.py | 622 device_assignment=None, argument
674 if device_assignment is None:
682 return device_assignment.host_device(replica=index)
685 return device_assignment.tpu_ordinal(replica=index)
746 device_assignment, argument
760 self._device_assignment = device_assignment
|
| D | tpu_estimator.py | 919 device_assignment=ctx.device_assignment)
974 if ctx.device_assignment:
975 return ctx.device_assignment.tpu_ordinal(replica=replica_id)
2489 device_assignment = ctx.device_assignment
2491 device_assignment = None
2495 tpu_computation, device_assignment=device_assignment)
2499 tpu_computation, device_assignment=device_assignment)
3185 device_assignment=ctx.device_assignment)
3210 device_assignment=ctx.device_assignment)
3240 device_assignment=ctx.device_assignment)
|
| D | device_assignment.py | 173 return device_assignment(topology, computation_shape, computation_stride,
177 def device_assignment(topology, function
|
| /external/tensorflow/tensorflow/compiler/xla/ |
| D | executable_run_options.cc | 80 const DeviceAssignment* device_assignment) { in set_device_assignment() argument
81 device_assignment_ = device_assignment; in set_device_assignment()
85 const DeviceAssignment* ExecutableRunOptions::device_assignment() const { in device_assignment() function in xla::ExecutableRunOptions
|
| D | executable_run_options.h | 77 const DeviceAssignment* device_assignment);
78 const DeviceAssignment* device_assignment() const;
|
| /external/tensorflow/tensorflow/compiler/xla/service/ |
| D | hlo_runner.cc | 274 DeviceAssignment* device_assignment, bool use_threads) { in ExecuteReplicated() argument
294 int64 device = (*device_assignment)(i, 0); in ExecuteReplicated()
300 device, streams.back().get(), device_assignment)); in ExecuteReplicated()
330 int64 device = (*device_assignment)(i, 0); in ExecuteReplicated()
348 int64 device = (*device_assignment)(i, 0); in ExecuteReplicated()
420 DeviceAssignment device_assignment, in ExecuteReplicated()
422 return ExecuteReplicated(std::move(module), options, &device_assignment, in ExecuteReplicated()
443 int64 device, se::Stream* stream, DeviceAssignment* device_assignment) { in GetServiceRunOptionsForDevice() argument
450 if (device_assignment != nullptr) { in GetServiceRunOptionsForDevice()
451 run_options.set_device_assignment(device_assignment); in GetServiceRunOptionsForDevice()
|
| D | hlo_module_config.h | 115 void set_static_device_assignment(const DeviceAssignment& device_assignment) { in set_static_device_assignment() argument
116 static_device_assignment_ = device_assignment; in set_static_device_assignment()
|
| D | hlo_runner.h | 180 DeviceAssignment* device_assignment, bool use_threads = false);
196 int64 device, se::Stream* stream, DeviceAssignment* device_assignment);
|
| D | compiler.h | 93 void set_static_device_assignment(const DeviceAssignment& device_assignment) { in set_static_device_assignment() argument
94 static_device_assignment_ = device_assignment; in set_static_device_assignment()
|
| D | service.cc | 327 auto device_assignment, in CreateModuleConfig()
328 DeviceAssignment::Deserialize(execution_options->device_assignment())); in CreateModuleConfig()
329 config->set_static_device_assignment(*device_assignment); in CreateModuleConfig()
412 DeviceAssignment device_assignment(options_.number_of_replicas(), in ExecuteParallelAndRegisterResult() local
418 device_assignment(replica, i) = replicas[replica]->device_ordinal(); in ExecuteParallelAndRegisterResult()
453 options.set_device_assignment(&device_assignment); in ExecuteParallelAndRegisterResult()
533 DeviceAssignment device_assignment(options_.number_of_replicas(), in ExecuteAndRegisterResult() local
536 device_assignment(replica, 0) = replicas[replica]->device_ordinal(); in ExecuteAndRegisterResult()
548 options.set_device_assignment(&device_assignment); in ExecuteAndRegisterResult()
|
| /external/tensorflow/tensorflow/contrib/distribute/python/ |
| D | combinations.py | 63 from tensorflow.python.tpu import device_assignment as device_assignment_lib
334 device_assignment = None
336 device_assignment = device_assignment_lib.DeviceAssignment(
341 device_assignment=device_assignment,
|
| /external/tensorflow/tensorflow/python/distribute/ |
| D | tpu_strategy.py | 43 from tensorflow.python.tpu import device_assignment as device_assignment_lib
131 device_assignment=None): argument
147 self, tpu_cluster_resolver, steps_per_run, device_assignment))
199 device_assignment=None): argument
212 self._device_assignment = device_assignment
|
| /external/tensorflow/tensorflow/contrib/tpu/python/tpu/ |
| D | device_assignment.py | 22 from tensorflow.python.tpu.device_assignment import *
|
| /external/tensorflow/tensorflow/contrib/tpu/ |
| D | __init__.py | 72 from tensorflow.contrib.tpu.python.tpu.device_assignment import *
|
| D | BUILD | 150 "python/tpu/device_assignment.py",
|
| /external/tensorflow/tensorflow/core/api_def/base_api/ |
| D | api_def_TPUReplicateMetadata.pbtxt | 29 name: "device_assignment"
|
| D | api_def_TPUReplicate.pbtxt | 66 name: "device_assignment"
|
| /external/tensorflow/tensorflow/tools/api/golden/v2/ |
| D | tensorflow.tpu.experimental.-device-assignment.pbtxt | 3 is_instance: "<class \'tensorflow.python.tpu.device_assignment.DeviceAssignment\'>"
|
| D | tensorflow.distribute.experimental.-t-p-u-strategy.pbtxt | 20 …argspec: "args=[\'self\', \'tpu_cluster_resolver\', \'steps_per_run\', \'device_assignment\'], var…
|
| /external/tensorflow/tensorflow/tools/api/golden/v1/ |
| D | tensorflow.tpu.experimental.-device-assignment.pbtxt | 3 is_instance: "<class \'tensorflow.python.tpu.device_assignment.DeviceAssignment\'>"
|
| D | tensorflow.distribute.experimental.-t-p-u-strategy.pbtxt | 20 …argspec: "args=[\'self\', \'tpu_cluster_resolver\', \'steps_per_run\', \'device_assignment\'], var…
|
| /external/tensorflow/tensorflow/compiler/xla/python/ |
| D | local_computation_builder.cc | 212 xla::DeviceAssignment device_assignment, xla::LocalClient* client) in LocalExecutable() argument
214 device_assignment_(std::move(device_assignment)), in LocalExecutable()
356 DeviceAssignment device_assignment, in Compile()
361 std::move(device_assignment), client.client()); in Compile()
|
| /external/tensorflow/tensorflow/compiler/xla/tests/ |
| D | hlo_test_base.cc | 222 int64 num_replicas, DeviceAssignment* device_assignment, in ExecuteReplicated() argument
231 device_assignment, use_threads); in ExecuteReplicated()
|