Home
last modified time | relevance | path

Searched defs:xla_device_assignment (Results 1 – 4 of 4) sorted by relevance

/aosp_15_r20/external/tensorflow/tensorflow/compiler/mlir/tensorflow/utils/
H A Dtpu_rewrite_device_util_test.cc106 llvm::Optional<xla::DeviceAssignmentProto> xla_device_assignment; in TEST_P() local
475 auto& xla_device_assignment = tpu_device_assignment.xla_device_assignment; in TEST() local
581 auto& xla_device_assignment = tpu_device_assignment.xla_device_assignment; in TEST() local
H A Dtpu_rewrite_device_util.h78 llvm::Optional<xla::DeviceAssignmentProto> xla_device_assignment; member
/aosp_15_r20/external/tensorflow/tensorflow/compiler/mlir/tensorflow/transforms/
H A Dtpu_rewrite_pass.cc281 llvm::Optional<xla::DeviceAssignmentProto>&& xla_device_assignment, in SetMetadataProtoFromClusterFuncOp()
334 llvm::Optional<xla::DeviceAssignmentProto>&& xla_device_assignment, in BuildCompileOp()
/aosp_15_r20/external/tensorflow/tensorflow/core/tpu/graph_rewrite/
H A Ddistributed_tpu_rewrite_pass.cc1742 std::unique_ptr<xla::DeviceAssignment>* xla_device_assignment) { in BuildGeneralDeviceAssignment()
1782 std::unique_ptr<xla::DeviceAssignment>* xla_device_assignment) { in BuildDeviceAssignment()
2604 const xla::DeviceAssignment* xla_device_assignment, in BuildCompileNode()
4376 std::unique_ptr<xla::DeviceAssignment>* xla_device_assignment, in GetDeviceTopology()
4730 std::unique_ptr<xla::DeviceAssignment> xla_device_assignment; in RewriteTPUReplicateNode() local