Searched defs:device_assignment_attr (Results 1 – 4 of 4) sorted by relevance
| /aosp_15_r20/external/tensorflow/tensorflow/compiler/mlir/tensorflow/utils/ |
| H A D | tpu_rewrite_device_util_test.cc | 59 std::vector<int64_t> device_assignment_attr; in TEST_P() local
300 std::vector<int64_t> device_assignment_attr{1, 0, 0, 0}; in TEST() local
323 std::vector<int64_t> device_assignment_attr; in TEST() local
419 std::vector<int64_t> device_assignment_attr{0, 0, 0, 0, 0, 0, 0, 1, 0, 1, 0, in TEST() local
533 std::vector<int64_t> device_assignment_attr{ in TEST() local
607 auto device_assignment_attr = builder.getI64ArrayAttr({1, 2, 3}); in TEST() local
620 auto device_assignment_attr = builder.getF32ArrayAttr({1.0, 2.0, 3.0}); in TEST() local
|
| H A D | tpu_rewrite_device_util.cc | 346 llvm::ArrayRef<int64_t> device_assignment_attr) { in GetGeneralTPUExecutionDeviceAssignment()
425 mlir::ArrayAttr device_assignment_attr) { in GetDeviceCoordinates()
448 llvm::ArrayRef<int64_t> device_assignment_attr) { in GetTPUCompilationAndExecutionDevices()
515 auto device_assignment_attr = cluster->getAttrOfType<mlir::ArrayAttr>( in GetHostDeviceOutsideComputation() local
|
| /aosp_15_r20/external/tensorflow/tensorflow/compiler/mlir/tensorflow/transforms/ |
| H A D | tpu_rewrite_pass.cc | 738 auto device_assignment_attr = cluster_func->getAttrOfType<mlir::ArrayAttr>( in Rewrite() local
|
| /aosp_15_r20/external/tensorflow/tensorflow/core/tpu/graph_rewrite/ |
| H A D | distributed_tpu_rewrite_pass.cc | 1652 absl::Span<const int> device_assignment_attr, in ParseDeviceAssignmentAttr()
1779 absl::Span<const int> device_assignment_attr, in BuildDeviceAssignment()
|