Searched defs:tpu_devices (Results 1 – 8 of 8) sorted by relevance
| /aosp_15_r20/external/tensorflow/tensorflow/compiler/mlir/tensorflow/utils/ |
| H A D | tpu_rewrite_device_util.cc | 119 llvm::SmallVectorImpl<llvm::SmallVector<Device, 8>>* tpu_devices) { in GetTPUDevices()
185 llvm::ArrayRef<llvm::SmallVector<Device, 8>> tpu_devices) { in GetFullMeshTPUExecutionDeviceAssignment()
344 llvm::ArrayRef<llvm::SmallVector<Device, 8>> tpu_devices, in GetGeneralTPUExecutionDeviceAssignment()
454 llvm::SmallVector<llvm::SmallVector<Device, 8>, 8> tpu_devices; in GetTPUCompilationAndExecutionDevices() local
|
| H A D | tpu_rewrite_device_util_test.cc | 334 const auto& tpu_devices = tpu_device_assignment.tpu_devices; in TEST() local
437 const auto& tpu_devices = tpu_device_assignment.tpu_devices; in TEST() local
551 auto& tpu_devices = tpu_device_assignment.tpu_devices; in TEST() local
|
| H A D | tpu_rewrite_device_util.h | 77 TPUDevicesAndHosts tpu_devices; member
|
| /aosp_15_r20/external/tensorflow/tensorflow/compiler/mlir/tensorflow/transforms/ |
| H A D | tpu_rewrite_pass.cc | 404 tpu_devices, in AssignDevicesToReplicate()
507 tpu_devices, in AddToParallelExecuteOp()
588 tpu_devices, in AssignDevicesToReplicatedExecute()
|
| /aosp_15_r20/external/tensorflow/tensorflow/core/tpu/graph_rewrite/ |
| H A D | distributed_tpu_rewrite_helpers.cc | 152 int* num_tpus_per_host, std::vector<std::vector<Device*>>* tpu_devices) { in GetTPUDevices()
|
| H A D | update_tpu_embedding_ops_passes.cc | 55 std::vector<Device*> tpu_devices; in Run() local
|
| H A D | distributed_tpu_rewrite_pass.cc | 1569 std::vector<std::vector<Device*>>* tpu_devices) { in GetTPUDeviceNames()
1713 int num_replicas, const std::vector<std::vector<Device*>>& tpu_devices, in BuildFullMeshDeviceAssignment()
1737 const std::vector<std::vector<Device*>>& tpu_devices, in BuildGeneralDeviceAssignment()
1777 const std::vector<std::vector<Device*>>& tpu_devices, int num_replicas, in BuildDeviceAssignment()
4387 std::vector<std::vector<Device*>> tpu_devices; in GetDeviceTopology() local
|
| /aosp_15_r20/external/tensorflow/tensorflow/core/tpu/kernels/ |
| H A D | tpu_functional_ops.cc | 1219 std::vector<Device*> tpu_devices; in ComputeAsync() local
1457 const std::vector<string>& tpu_devices) { in InitializeShardedVarOnTPU()
1861 std::vector<string> tpu_devices; in ReplaceAndPartitionXLAShardingVariable() local
|