Home
last modified time | relevance | path

Searched refs:tpu_device_assignment (Results 1 – 3 of 3) sorted by relevance

/external/tensorflow/tensorflow/compiler/mlir/tensorflow/utils/
Dtpu_rewrite_device_util_test.cc330 const auto& tpu_device_assignment = status_or.ValueOrDie(); in TEST() local
331 EXPECT_EQ(tpu_device_assignment.compilation_device, in TEST()
333 const auto& tpu_devices = tpu_device_assignment.tpu_devices; in TEST()
371 EXPECT_FALSE(tpu_device_assignment.xla_device_assignment.hasValue()); in TEST()
433 const auto& tpu_device_assignment = status_or.ValueOrDie(); in TEST() local
434 EXPECT_EQ(tpu_device_assignment.compilation_device, in TEST()
436 const auto& tpu_devices = tpu_device_assignment.tpu_devices; in TEST()
474 auto& xla_device_assignment = tpu_device_assignment.xla_device_assignment; in TEST()
546 auto& tpu_device_assignment = status_or.ValueOrDie(); in TEST() local
547 EXPECT_EQ(tpu_device_assignment.compilation_device, in TEST()
[all …]
Dtpu_rewrite_device_util.cc538 auto& tpu_device_assignment = status_or_tpu_device_assignment.ValueOrDie(); in GetHostDeviceOutsideComputation() local
540 *host_device = tpu_device_assignment.tpu_devices[0][0].host; in GetHostDeviceOutsideComputation()
/external/tensorflow/tensorflow/compiler/mlir/tensorflow/transforms/
Dtpu_rewrite_pass.cc617 auto& tpu_device_assignment = status_or_tpu_device_assignment.ValueOrDie(); in Rewrite() local
631 tpu_device_assignment.compilation_device, in Rewrite()
632 std::move(tpu_device_assignment.xla_device_assignment), builder); in Rewrite()
657 compile_op, tpu_device_assignment.compilation_device, builder); in Rewrite()
659 AssignDevicesToReplicate(replicate, tpu_device_assignment.tpu_devices, in Rewrite()
673 result = BuildParallelExecuteOp(tpu_device_assignment.tpu_devices, in Rewrite()
697 tpu_device_assignment.tpu_devices, execute_op, builder); in Rewrite()