Searched defs:xla_device_assignment (Results 1 – 4 of 4) sorted by relevance
76 llvm::Optional<xla::DeviceAssignmentProto> xla_device_assignment; member
105 llvm::Optional<xla::DeviceAssignmentProto> xla_device_assignment; in TEST_P() local474 auto& xla_device_assignment = tpu_device_assignment.xla_device_assignment; in TEST() local580 auto& xla_device_assignment = tpu_device_assignment.xla_device_assignment; in TEST() local
304 llvm::Optional<xla::DeviceAssignmentProto>&& xla_device_assignment, in SetMetadataProtoFromClusterFuncOp()352 llvm::Optional<xla::DeviceAssignmentProto>&& xla_device_assignment, in BuildCompileOp()
1529 std::unique_ptr<xla::DeviceAssignment>* xla_device_assignment) { in BuildGeneralDeviceAssignment()1567 std::unique_ptr<xla::DeviceAssignment>* xla_device_assignment) { in BuildDeviceAssignment()2316 const xla::DeviceAssignment* xla_device_assignment, in BuildCompileNode()3921 std::unique_ptr<xla::DeviceAssignment>* xla_device_assignment, in GetDeviceTopology()4228 std::unique_ptr<xla::DeviceAssignment> xla_device_assignment; in RewriteTPUReplicateNode() local