Searched defs:device_assignment_attr (Results 1 – 4 of 4) sorted by relevance
| /external/tensorflow/tensorflow/compiler/mlir/tensorflow/utils/ |
| D | tpu_rewrite_device_util_test.cc | 58 std::vector<int64_t> device_assignment_attr; in TEST_P() local
299 std::vector<int64_t> device_assignment_attr{1, 0, 0, 0}; in TEST() local
322 std::vector<int64_t> device_assignment_attr; in TEST() local
418 std::vector<int64_t> device_assignment_attr{0, 0, 0, 0, 0, 0, 0, 1, 0, 1, 0, in TEST() local
532 std::vector<int64_t> device_assignment_attr{ in TEST() local
606 auto device_assignment_attr = builder.getI64ArrayAttr({1, 2, 3}); in TEST() local
619 auto device_assignment_attr = builder.getF32ArrayAttr({1.0, 2.0, 3.0}); in TEST() local
|
| D | tpu_rewrite_device_util.cc | 350 llvm::ArrayRef<int64_t> device_assignment_attr) { in GetGeneralTPUExecutionDeviceAssignment()
429 mlir::ArrayAttr device_assignment_attr) { in GetDeviceCoordinates()
452 llvm::ArrayRef<int64_t> device_assignment_attr) { in GetTPUCompilationAndExecutionDevices()
513 auto device_assignment_attr = cluster->getAttrOfType<mlir::ArrayAttr>( in GetHostDeviceOutsideComputation() local
|
| /external/tensorflow/tensorflow/compiler/mlir/tensorflow/transforms/ |
| D | tpu_rewrite_pass.cc | 590 auto device_assignment_attr = cluster_func->getAttrOfType<mlir::ArrayAttr>( in Rewrite() local
|
| /external/tensorflow/tensorflow/core/tpu/graph_rewrite/ |
| D | distributed_tpu_rewrite_pass.cc | 1442 absl::Span<const int> device_assignment_attr, in ParseDeviceAssignmentAttr()
1565 absl::Span<const int> device_assignment_attr, in BuildDeviceAssignment()
|