Searched defs:tpu_device (Results 1 – 10 of 10) sorted by relevance
76 string tpu_device; in ComputeAsync() local233 std::string tpu_device; in Compute() local
1874 string tpu_device = strings::StrCat(kTPUDeviceNamePrefix, device_index); in ReplaceAndPartitionXLAShardingVariable() local
172 std::string GetCPUHostDeviceForTPUDevice(Device tpu_device) { in GetCPUHostDeviceForTPUDevice()204 const auto& tpu_device = tpu_devices[task][device]; in GetFullMeshTPUExecutionDeviceAssignment() local409 const auto& tpu_device = tpu_devices[task][device]; in GetGeneralTPUExecutionDeviceAssignment() local
473 def generate_dequeue_op(self, tpu_device=0): argument788 def generate_dequeue_op(self, tpu_device=0): argument
165 def tpu_device(self, member in DeviceAssignment
75 llvm::StringRef tpu_device = tpu_replica_0.cast<ArrayAttr>()[replica_id] in GetDeviceOrdinal() local
88 for (const auto* tpu_device : tpu_devices) { in Run() local
3920 const DeviceNameUtils::ParsedName& tpu_device, in CopyOutsideCompilationNodes()3978 DeviceNameUtils::ParsedName tpu_device; in ReplicateOutsideCompilationNodes() local
92 DeviceProperties tpu_device; in TEST() local
331 py::class_<PjRtTpuDevice, PjRtDevice, ClientAndPtr<PjRtTpuDevice>> tpu_device( in PYBIND11_MODULE() local