Searched refs:xla_device_assignment (Results 1 – 5 of 5) sorted by relevance
| /external/tensorflow/tensorflow/compiler/mlir/tensorflow/utils/ |
| D | tpu_rewrite_device_util.h | 71 xla::DeviceAssignmentProto&& xla_device_assignment) in TPUDeviceAssignment()
74 xla_device_assignment(std::move(xla_device_assignment)) {} in TPUDeviceAssignment()
78 llvm::Optional<xla::DeviceAssignmentProto> xla_device_assignment; member
|
| D | tpu_rewrite_device_util_test.cc | 106 llvm::Optional<xla::DeviceAssignmentProto> xla_device_assignment; in TEST_P() local
372 EXPECT_FALSE(tpu_device_assignment.xla_device_assignment.has_value()); in TEST()
475 auto& xla_device_assignment = tpu_device_assignment.xla_device_assignment; in TEST() local
476 ASSERT_TRUE(xla_device_assignment.has_value()); in TEST()
477 EXPECT_EQ(xla_device_assignment->replica_count(), 4); in TEST()
478 EXPECT_EQ(xla_device_assignment->computation_count(), 2); in TEST()
479 ASSERT_EQ(xla_device_assignment->computation_devices_size(), 2); in TEST()
481 xla_device_assignment->computation_devices(0); in TEST()
484 xla_device_assignment->computation_devices(1); in TEST()
581 auto& xla_device_assignment = tpu_device_assignment.xla_device_assignment; in TEST() local
[all …]
|
| /external/tensorflow/tensorflow/compiler/mlir/tensorflow/transforms/ |
| D | tpu_rewrite_pass.cc | 281 llvm::Optional<xla::DeviceAssignmentProto>&& xla_device_assignment, in SetMetadataProtoFromClusterFuncOp() argument
296 if (xla_device_assignment.has_value()) in SetMetadataProtoFromClusterFuncOp()
298 std::move(xla_device_assignment.getValue()); in SetMetadataProtoFromClusterFuncOp()
334 llvm::Optional<xla::DeviceAssignmentProto>&& xla_device_assignment, in BuildCompileOp() argument
340 std::move(xla_device_assignment), &metadata))) in BuildCompileOp()
772 std::move(tpu_device_assignment.xla_device_assignment), in Rewrite()
|
| /external/tensorflow/tensorflow/core/tpu/graph_rewrite/ |
| D | distributed_tpu_rewrite_pass.h | 281 std::unique_ptr<xla::DeviceAssignment>* xla_device_assignment);
373 const xla::DeviceAssignment* xla_device_assignment,
537 std::unique_ptr<xla::DeviceAssignment>* xla_device_assignment,
|
| D | distributed_tpu_rewrite_pass.cc | 1742 std::unique_ptr<xla::DeviceAssignment>* xla_device_assignment) { in BuildGeneralDeviceAssignment() argument
1745 *xla_device_assignment = absl::make_unique<xla::DeviceAssignment>( in BuildGeneralDeviceAssignment()
1761 (**xla_device_assignment)(replica, computation) = core_location.Id(); in BuildGeneralDeviceAssignment()
1782 std::unique_ptr<xla::DeviceAssignment>* xla_device_assignment) { in BuildDeviceAssignment() argument
1871 topology, tf_device_assignment, devices_to_lock, xla_device_assignment); in BuildDeviceAssignment()
2604 const xla::DeviceAssignment* xla_device_assignment, in BuildCompileNode() argument
2633 if (xla_device_assignment != nullptr) { in BuildCompileNode()
2635 xla_device_assignment->Serialize(proto.mutable_device_assignment())); in BuildCompileNode()
4376 std::unique_ptr<xla::DeviceAssignment>* xla_device_assignment, in GetDeviceTopology() argument
4421 devices_to_lock, xla_device_assignment)); in GetDeviceTopology()
[all …]
|